Merge branch 'main' into feat/litellm_sambanova_usage

.cursor/rules/general.mdc

@@ -1,9 +0,0 @@
----
-description: General rules always applicable across the project
-globs:
-alwaysApply: true
----
-# Style
-
-- Comments must add value to code. Don't write filler comments explaining what you are doing next; they just add noise.
-- Add a comment to clarify surprising behavior which would not be obvious. Good variable naming and clear code organization is more important.

.github/dependabot.yml

@@ -5,4 +5,19 @@ updates:
   - package-ecosystem: "github-actions"
     directory: "/" # Will use the default workflow location of `.github/workflows`
     schedule:
-      interval: "daily"
+      interval: "weekly"
+      day: "saturday"
+    commit-message:
+      prefix: chore(github-deps)
+  - package-ecosystem: "uv"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+      day: "saturday"
+    # ignore all non-security updates: https://docs.github.com/en/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file#open-pull-requests-limit
+    open-pull-requests-limit: 0
+    labels:
+      - type/dependencies
+      - python
+    commit-message:
+      prefix: chore(python-deps)

.github/workflows/integration-tests.yml

@@ -1,13 +1,28 @@
 name: Integration tests
 
 on:
-  pull_request:
   push:
-    branches: [main]
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+    paths:
+      - 'distributions/**'
+      - 'llama_stack/**'
+      - 'tests/integration/**'
+      - 'uv.lock'
+      - 'pyproject.toml'
+      - 'requirements.txt'
+      - '.github/workflows/integration-tests.yml' # This workflow
 
 jobs:
-  ollama:
+  test-matrix:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        # Listing tests manually since some of them currently fail
+        # TODO: generate matrix list from tests/integration when fixed
+        test-type: [inference, datasets, inspect, scoring, post_training, providers]
+      fail-fast: false # we want to run all tests regardless of failure
 
     steps:
       - name: Checkout repository
@@ -34,6 +49,8 @@ jobs:
         run: |
           uv sync --extra dev --extra test
           uv pip install ollama faiss-cpu
+          # always test against the latest version of the client
+          uv pip install git+https://github.com/meta-llama/llama-stack-client-python.git@main
           uv pip install -e .
 
       - name: Wait for Ollama to start
@@ -56,25 +73,24 @@ jobs:
           INFERENCE_MODEL: "meta-llama/Llama-3.2-3B-Instruct"
         run: |
           source .venv/bin/activate
-          # TODO: use "llama stack run"
+          nohup uv run llama stack run ./llama_stack/templates/ollama/run.yaml --image-type venv > server.log 2>&1 &
-          nohup uv run python -m llama_stack.distribution.server.server --yaml-config ./llama_stack/templates/ollama/run.yaml > server.log 2>&1 &
 
       - name: Wait for Llama Stack server to be ready
         run: |
           echo "Waiting for Llama Stack server..."
           for i in {1..30}; do
             if curl -s http://localhost:8321/v1/health | grep -q "OK"; then
-              echo " Llama Stack server is up!"
+              echo "Llama Stack server is up!"
               exit 0
             fi
             sleep 1
           done
-          echo " Llama Stack server failed to start"
+          echo "Llama Stack server failed to start"
           cat server.log
           exit 1
 
-      - name: Run Inference Integration Tests
+      - name: Run Integration Tests
         env:
           INFERENCE_MODEL: "meta-llama/Llama-3.2-3B-Instruct"
         run: |
-          uv run pytest -v tests/integration/inference --stack-config=ollama --text-model="meta-llama/Llama-3.2-3B-Instruct" --embedding-model=all-MiniLM-L6-v2
+          uv run pytest -v tests/integration/${{ matrix.test-type }} --stack-config=ollama --text-model="meta-llama/Llama-3.2-3B-Instruct" --embedding-model=all-MiniLM-L6-v2

.github/workflows/providers-build.yml

@@ -40,6 +40,7 @@ jobs:
       matrix:
         template: ${{ fromJson(needs.generate-matrix.outputs.templates) }}
         image-type: [venv, container]
+      fail-fast: false # We want to run all jobs even if some fail
 
     steps:
       - name: Checkout repository
@@ -67,7 +68,9 @@ jobs:
 
       - name: Run Llama Stack Build
         run: |
-          uv run llama stack build --template ${{ matrix.template }} --image-type ${{ matrix.image-type }} --image-name test
+          # USE_COPY_NOT_MOUNT is set to true since mounting is not supported by docker buildx, we use COPY instead
+          # LLAMA_STACK_DIR is set to the current directory so we are building from the source
+          USE_COPY_NOT_MOUNT=true LLAMA_STACK_DIR=. uv run llama stack build --template ${{ matrix.template }} --image-type ${{ matrix.image-type }} --image-name test
 
       - name: Print dependencies in the image
         if: matrix.image-type == 'venv'

.github/workflows/unit-tests.yml

@@ -5,6 +5,14 @@ on:
     branches: [ main ]
   pull_request:
     branches: [ main ]
+    paths:
+      - 'distributions/**'
+      - 'llama_stack/**'
+      - 'tests/unit/**'
+      - 'uv.lock'
+      - 'pyproject.toml'
+      - 'requirements.txt'
+      - '.github/workflows/unit-tests.yml' # This workflow
   workflow_dispatch:
 
 jobs:

.pre-commit-config.yaml

@@ -77,7 +77,7 @@ repos:
         name: Distribution Template Codegen
         additional_dependencies:
           - uv==0.6.0
-        entry: uv run --extra codegen python -m llama_stack.scripts.distro_codegen
+        entry: uv run --extra codegen ./scripts/distro_codegen.py
         language: python
         pass_filenames: false
         require_serial: true

CONTRIBUTING.md

@@ -159,7 +159,7 @@ LLAMA_STACK_DIR=$(pwd) LLAMA_STACK_CLIENT_DIR=../llama-stack-client-python llama
 
 ### Updating Provider Configurations
 
-If you have made changes to a provider's configuration in any form (introducing a new config key, or changing models, etc.), you should run `python llama_stack/scripts/distro_codegen.py` to re-generate various YAML files as well as the documentation. You should not change `docs/source/.../distributions/` files manually as they are auto-generated.
+If you have made changes to a provider's configuration in any form (introducing a new config key, or changing models, etc.), you should run `./scripts/distro_codegen.py` to re-generate various YAML files as well as the documentation. You should not change `docs/source/.../distributions/` files manually as they are auto-generated.
 
 ### Building the Documentation
 

distributions/dependencies.json

@@ -401,16 +401,13 @@
   ],
   "nvidia": [
     "aiosqlite",
-    "autoevals",
     "blobfile",
     "chardet",
-    "datasets",
     "faiss-cpu",
     "fastapi",
     "fire",
     "httpx",
     "matplotlib",
-    "mcp",
     "nltk",
     "numpy",
     "openai",

docs/_static/llama-stack-spec.html

@@ -40,75 +40,7 @@
         }
     ],
     "paths": {
-        "/v1/datasetio/rows": {
+        "/v1/datasetio/append-rows/{dataset_id}": {
-            "get": {
-                "responses": {
-                    "200": {
-                        "description": "OK",
-                        "content": {
-                            "application/json": {
-                                "schema": {
-                                    "$ref": "#/components/schemas/PaginatedRowsResult"
-                                }
-                            }
-                        }
-                    },
-                    "400": {
-                        "$ref": "#/components/responses/BadRequest400"
-                    },
-                    "429": {
-                        "$ref": "#/components/responses/TooManyRequests429"
-                    },
-                    "500": {
-                        "$ref": "#/components/responses/InternalServerError500"
-                    },
-                    "default": {
-                        "$ref": "#/components/responses/DefaultError"
-                    }
-                },
-                "tags": [
-                    "DatasetIO"
-                ],
-                "description": "Get a paginated list of rows from a dataset.",
-                "parameters": [
-                    {
-                        "name": "dataset_id",
-                        "in": "query",
-                        "description": "The ID of the dataset to get the rows from.",
-                        "required": true,
-                        "schema": {
-                            "type": "string"
-                        }
-                    },
-                    {
-                        "name": "rows_in_page",
-                        "in": "query",
-                        "description": "The number of rows to get per page.",
-                        "required": true,
-                        "schema": {
-                            "type": "integer"
-                        }
-                    },
-                    {
-                        "name": "page_token",
-                        "in": "query",
-                        "description": "The token to get the next page of rows.",
-                        "required": false,
-                        "schema": {
-                            "type": "string"
-                        }
-                    },
-                    {
-                        "name": "filter_condition",
-                        "in": "query",
-                        "description": "(Optional) A condition to filter the rows by.",
-                        "required": false,
-                        "schema": {
-                            "type": "string"
-                        }
-                    }
-                ]
-            },
             "post": {
                 "responses": {
                     "200": {
@@ -131,7 +63,16 @@
                     "DatasetIO"
                 ],
                 "description": "",
-                "parameters": [],
+                "parameters": [
+                    {
+                        "name": "dataset_id",
+                        "in": "path",
+                        "required": true,
+                        "schema": {
+                            "type": "string"
+                        }
+                    }
+                ],
                 "requestBody": {
                     "content": {
                         "application/json": {
@@ -583,7 +524,7 @@
                     }
                 },
                 "tags": [
-                    "Files (Coming Soon)"
+                    "Files"
                 ],
                 "description": "List all buckets.",
                 "parameters": [
@@ -623,7 +564,7 @@
                     }
                 },
                 "tags": [
-                    "Files (Coming Soon)"
+                    "Files"
                 ],
                 "description": "Create a new upload session for a file identified by a bucket and key.",
                 "parameters": [],
@@ -850,7 +791,7 @@
                     }
                 },
                 "tags": [
-                    "Files (Coming Soon)"
+                    "Files"
                 ],
                 "description": "Get a file info identified by a bucket and key.",
                 "parameters": [
@@ -900,7 +841,7 @@
                     }
                 },
                 "tags": [
-                    "Files (Coming Soon)"
+                    "Files"
                 ],
                 "description": "Delete a file identified by a bucket and key.",
                 "parameters": [
@@ -1889,7 +1830,7 @@
                     }
                 },
                 "tags": [
-                    "Files (Coming Soon)"
+                    "Files"
                 ],
                 "description": "Returns information about an existsing upload session",
                 "parameters": [
@@ -1937,7 +1878,7 @@
                     }
                 },
                 "tags": [
-                    "Files (Coming Soon)"
+                    "Files"
                 ],
                 "description": "Upload file content to an existing upload session. On the server, request body will have the raw bytes that are uploaded.",
                 "parameters": [
@@ -2236,6 +2177,67 @@
                 }
             }
         },
+        "/v1/datasetio/iterrows/{dataset_id}": {
+            "get": {
+                "responses": {
+                    "200": {
+                        "description": "OK",
+                        "content": {
+                            "application/json": {
+                                "schema": {
+                                    "$ref": "#/components/schemas/IterrowsResponse"
+                                }
+                            }
+                        }
+                    },
+                    "400": {
+                        "$ref": "#/components/responses/BadRequest400"
+                    },
+                    "429": {
+                        "$ref": "#/components/responses/TooManyRequests429"
+                    },
+                    "500": {
+                        "$ref": "#/components/responses/InternalServerError500"
+                    },
+                    "default": {
+                        "$ref": "#/components/responses/DefaultError"
+                    }
+                },
+                "tags": [
+                    "DatasetIO"
+                ],
+                "description": "Get a paginated list of rows from a dataset. Uses cursor-based pagination.",
+                "parameters": [
+                    {
+                        "name": "dataset_id",
+                        "in": "path",
+                        "description": "The ID of the dataset to get the rows from.",
+                        "required": true,
+                        "schema": {
+                            "type": "string"
+                        }
+                    },
+                    {
+                        "name": "start_index",
+                        "in": "query",
+                        "description": "Index into dataset for the first row to get. Get all rows if None.",
+                        "required": false,
+                        "schema": {
+                            "type": "integer"
+                        }
+                    },
+                    {
+                        "name": "limit",
+                        "in": "query",
+                        "description": "The number of rows to get.",
+                        "required": false,
+                        "schema": {
+                            "type": "integer"
+                        }
+                    }
+                ]
+            }
+        },
         "/v1/eval/benchmarks/{benchmark_id}/jobs/{job_id}": {
             "get": {
                 "responses": {
@@ -2535,7 +2537,14 @@
             "post": {
                 "responses": {
                     "200": {
-                        "description": "OK"
+                        "description": "OK",
+                        "content": {
+                            "application/json": {
+                                "schema": {
+                                    "$ref": "#/components/schemas/Dataset"
+                                }
+                            }
+                        }
                     },
                     "400": {
                         "$ref": "#/components/responses/BadRequest400"
@@ -2553,7 +2562,7 @@
                 "tags": [
                     "Datasets"
                 ],
-                "description": "",
+                "description": "Register a new dataset.",
                 "parameters": [],
                 "requestBody": {
                     "content": {
@@ -2594,7 +2603,7 @@
                     }
                 },
                 "tags": [
-                    "Files (Coming Soon)"
+                    "Files"
                 ],
                 "description": "List all files in a bucket.",
                 "parameters": [
@@ -3824,9 +3833,6 @@
             "AppendRowsRequest": {
                 "type": "object",
                 "properties": {
-                    "dataset_id": {
-                        "type": "string"
-                    },
                     "rows": {
                         "type": "array",
                         "items": {
@@ -3858,7 +3864,6 @@
                 },
                 "additionalProperties": false,
                 "required": [
-                    "dataset_id",
                     "rows"
                 ],
                 "title": "AppendRowsRequest"
@@ -6824,6 +6829,224 @@
                 ],
                 "title": "Benchmark"
             },
+            "DataSource": {
+                "oneOf": [
+                    {
+                        "$ref": "#/components/schemas/URIDataSource"
+                    },
+                    {
+                        "$ref": "#/components/schemas/RowsDataSource"
+                    }
+                ],
+                "discriminator": {
+                    "propertyName": "type",
+                    "mapping": {
+                        "uri": "#/components/schemas/URIDataSource",
+                        "rows": "#/components/schemas/RowsDataSource"
+                    }
+                }
+            },
+            "Dataset": {
+                "type": "object",
+                "properties": {
+                    "identifier": {
+                        "type": "string"
+                    },
+                    "provider_resource_id": {
+                        "type": "string"
+                    },
+                    "provider_id": {
+                        "type": "string"
+                    },
+                    "type": {
+                        "type": "string",
+                        "const": "dataset",
+                        "default": "dataset"
+                    },
+                    "purpose": {
+                        "type": "string",
+                        "enum": [
+                            "post-training/messages",
+                            "eval/question-answer",
+                            "eval/messages-answer"
+                        ],
+                        "title": "DatasetPurpose",
+                        "description": "Purpose of the dataset. Each purpose has a required input data schema."
+                    },
+                    "source": {
+                        "$ref": "#/components/schemas/DataSource"
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "additionalProperties": {
+                            "oneOf": [
+                                {
+                                    "type": "null"
+                                },
+                                {
+                                    "type": "boolean"
+                                },
+                                {
+                                    "type": "number"
+                                },
+                                {
+                                    "type": "string"
+                                },
+                                {
+                                    "type": "array"
+                                },
+                                {
+                                    "type": "object"
+                                }
+                            ]
+                        }
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "identifier",
+                    "provider_resource_id",
+                    "provider_id",
+                    "type",
+                    "purpose",
+                    "source",
+                    "metadata"
+                ],
+                "title": "Dataset"
+            },
+            "RowsDataSource": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "rows",
+                        "default": "rows"
+                    },
+                    "rows": {
+                        "type": "array",
+                        "items": {
+                            "type": "object",
+                            "additionalProperties": {
+                                "oneOf": [
+                                    {
+                                        "type": "null"
+                                    },
+                                    {
+                                        "type": "boolean"
+                                    },
+                                    {
+                                        "type": "number"
+                                    },
+                                    {
+                                        "type": "string"
+                                    },
+                                    {
+                                        "type": "array"
+                                    },
+                                    {
+                                        "type": "object"
+                                    }
+                                ]
+                            }
+                        },
+                        "description": "The dataset is stored in rows. E.g. - [ {\"messages\": [{\"role\": \"user\", \"content\": \"Hello, world!\"}, {\"role\": \"assistant\", \"content\": \"Hello, world!\"}]} ]"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type",
+                    "rows"
+                ],
+                "title": "RowsDataSource",
+                "description": "A dataset stored in rows."
+            },
+            "URIDataSource": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "uri",
+                        "default": "uri"
+                    },
+                    "uri": {
+                        "type": "string",
+                        "description": "The dataset can be obtained from a URI. E.g. - \"https://mywebsite.com/mydata.jsonl\" - \"lsfs://mydata.jsonl\" - \"data:csv;base64,{base64_content}\""
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type",
+                    "uri"
+                ],
+                "title": "URIDataSource",
+                "description": "A dataset that can be obtained from a URI."
+            },
+            "Model": {
+                "type": "object",
+                "properties": {
+                    "identifier": {
+                        "type": "string"
+                    },
+                    "provider_resource_id": {
+                        "type": "string"
+                    },
+                    "provider_id": {
+                        "type": "string"
+                    },
+                    "type": {
+                        "type": "string",
+                        "const": "model",
+                        "default": "model"
+                    },
+                    "metadata": {
+                        "type": "object",
+                        "additionalProperties": {
+                            "oneOf": [
+                                {
+                                    "type": "null"
+                                },
+                                {
+                                    "type": "boolean"
+                                },
+                                {
+                                    "type": "number"
+                                },
+                                {
+                                    "type": "string"
+                                },
+                                {
+                                    "type": "array"
+                                },
+                                {
+                                    "type": "object"
+                                }
+                            ]
+                        }
+                    },
+                    "model_type": {
+                        "$ref": "#/components/schemas/ModelType",
+                        "default": "llm"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "identifier",
+                    "provider_resource_id",
+                    "provider_id",
+                    "type",
+                    "metadata",
+                    "model_type"
+                ],
+                "title": "Model"
+            },
+            "ModelType": {
+                "type": "string",
+                "enum": [
+                    "llm",
+                    "embedding"
+                ],
+                "title": "ModelType"
+            },
             "AgentTurnInputType": {
                 "type": "object",
                 "properties": {
@@ -6899,70 +7122,6 @@
                 ],
                 "title": "CompletionInputType"
             },
-            "Dataset": {
-                "type": "object",
-                "properties": {
-                    "identifier": {
-                        "type": "string"
-                    },
-                    "provider_resource_id": {
-                        "type": "string"
-                    },
-                    "provider_id": {
-                        "type": "string"
-                    },
-                    "type": {
-                        "type": "string",
-                        "const": "dataset",
-                        "default": "dataset"
-                    },
-                    "dataset_schema": {
-                        "type": "object",
-                        "additionalProperties": {
-                            "$ref": "#/components/schemas/ParamType"
-                        }
-                    },
-                    "url": {
-                        "$ref": "#/components/schemas/URL"
-                    },
-                    "metadata": {
-                        "type": "object",
-                        "additionalProperties": {
-                            "oneOf": [
-                                {
-                                    "type": "null"
-                                },
-                                {
-                                    "type": "boolean"
-                                },
-                                {
-                                    "type": "number"
-                                },
-                                {
-                                    "type": "string"
-                                },
-                                {
-                                    "type": "array"
-                                },
-                                {
-                                    "type": "object"
-                                }
-                            ]
-                        }
-                    }
-                },
-                "additionalProperties": false,
-                "required": [
-                    "identifier",
-                    "provider_resource_id",
-                    "provider_id",
-                    "type",
-                    "dataset_schema",
-                    "url",
-                    "metadata"
-                ],
-                "title": "Dataset"
-            },
             "JsonType": {
                 "type": "object",
                 "properties": {
@@ -7057,151 +7216,6 @@
                     }
                 }
             },
-            "StringType": {
-                "type": "object",
-                "properties": {
-                    "type": {
-                        "type": "string",
-                        "const": "string",
-                        "default": "string"
-                    }
-                },
-                "additionalProperties": false,
-                "required": [
-                    "type"
-                ],
-                "title": "StringType"
-            },
-            "UnionType": {
-                "type": "object",
-                "properties": {
-                    "type": {
-                        "type": "string",
-                        "const": "union",
-                        "default": "union"
-                    }
-                },
-                "additionalProperties": false,
-                "required": [
-                    "type"
-                ],
-                "title": "UnionType"
-            },
-            "Model": {
-                "type": "object",
-                "properties": {
-                    "identifier": {
-                        "type": "string"
-                    },
-                    "provider_resource_id": {
-                        "type": "string"
-                    },
-                    "provider_id": {
-                        "type": "string"
-                    },
-                    "type": {
-                        "type": "string",
-                        "const": "model",
-                        "default": "model"
-                    },
-                    "metadata": {
-                        "type": "object",
-                        "additionalProperties": {
-                            "oneOf": [
-                                {
-                                    "type": "null"
-                                },
-                                {
-                                    "type": "boolean"
-                                },
-                                {
-                                    "type": "number"
-                                },
-                                {
-                                    "type": "string"
-                                },
-                                {
-                                    "type": "array"
-                                },
-                                {
-                                    "type": "object"
-                                }
-                            ]
-                        }
-                    },
-                    "model_type": {
-                        "$ref": "#/components/schemas/ModelType",
-                        "default": "llm"
-                    }
-                },
-                "additionalProperties": false,
-                "required": [
-                    "identifier",
-                    "provider_resource_id",
-                    "provider_id",
-                    "type",
-                    "metadata",
-                    "model_type"
-                ],
-                "title": "Model"
-            },
-            "ModelType": {
-                "type": "string",
-                "enum": [
-                    "llm",
-                    "embedding"
-                ],
-                "title": "ModelType"
-            },
-            "PaginatedRowsResult": {
-                "type": "object",
-                "properties": {
-                    "rows": {
-                        "type": "array",
-                        "items": {
-                            "type": "object",
-                            "additionalProperties": {
-                                "oneOf": [
-                                    {
-                                        "type": "null"
-                                    },
-                                    {
-                                        "type": "boolean"
-                                    },
-                                    {
-                                        "type": "number"
-                                    },
-                                    {
-                                        "type": "string"
-                                    },
-                                    {
-                                        "type": "array"
-                                    },
-                                    {
-                                        "type": "object"
-                                    }
-                                ]
-                            }
-                        },
-                        "description": "The rows in the current page."
-                    },
-                    "total_count": {
-                        "type": "integer",
-                        "description": "The total number of rows in the dataset."
-                    },
-                    "next_page_token": {
-                        "type": "string",
-                        "description": "The token to get the next page of rows."
-                    }
-                },
-                "additionalProperties": false,
-                "required": [
-                    "rows",
-                    "total_count"
-                ],
-                "title": "PaginatedRowsResult",
-                "description": "A paginated list of rows from a dataset."
-            },
             "ScoringFn": {
                 "type": "object",
                 "properties": {
@@ -7265,6 +7279,36 @@
                 ],
                 "title": "ScoringFn"
             },
+            "StringType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "string",
+                        "default": "string"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ],
+                "title": "StringType"
+            },
+            "UnionType": {
+                "type": "object",
+                "properties": {
+                    "type": {
+                        "type": "string",
+                        "const": "union",
+                        "default": "union"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "type"
+                ],
+                "title": "UnionType"
+            },
             "Shield": {
                 "type": "object",
                 "properties": {
@@ -8084,6 +8128,50 @@
                 ],
                 "title": "ToolInvocationResult"
             },
+            "IterrowsResponse": {
+                "type": "object",
+                "properties": {
+                    "data": {
+                        "type": "array",
+                        "items": {
+                            "type": "object",
+                            "additionalProperties": {
+                                "oneOf": [
+                                    {
+                                        "type": "null"
+                                    },
+                                    {
+                                        "type": "boolean"
+                                    },
+                                    {
+                                        "type": "number"
+                                    },
+                                    {
+                                        "type": "string"
+                                    },
+                                    {
+                                        "type": "array"
+                                    },
+                                    {
+                                        "type": "object"
+                                    }
+                                ]
+                            }
+                        },
+                        "description": "The rows in the current page."
+                    },
+                    "next_start_index": {
+                        "type": "integer",
+                        "description": "Index into dataset for the first row in the next page. None if there are no more rows."
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "data"
+                ],
+                "title": "IterrowsResponse",
+                "description": "A paginated list of rows from a dataset."
+            },
             "ListAgentSessionsResponse": {
                 "type": "object",
                 "properties": {
@@ -9330,23 +9418,18 @@
             "RegisterDatasetRequest": {
                 "type": "object",
                 "properties": {
-                    "dataset_id": {
+                    "purpose": {
-                        "type": "string"
+                        "type": "string",
+                        "enum": [
+                            "post-training/messages",
+                            "eval/question-answer",
+                            "eval/messages-answer"
+                        ],
+                        "description": "The purpose of the dataset. One of - \"post-training/messages\": The dataset contains a messages column with list of messages for post-training. { \"messages\": [ {\"role\": \"user\", \"content\": \"Hello, world!\"}, {\"role\": \"assistant\", \"content\": \"Hello, world!\"}, ] } - \"eval/question-answer\": The dataset contains a question column and an answer column for evaluation. { \"question\": \"What is the capital of France?\", \"answer\": \"Paris\" } - \"eval/messages-answer\": The dataset contains a messages column with list of messages and an answer column for evaluation. { \"messages\": [ {\"role\": \"user\", \"content\": \"Hello, my name is John Doe.\"}, {\"role\": \"assistant\", \"content\": \"Hello, John Doe. How can I help you today?\"}, {\"role\": \"user\", \"content\": \"What's my name?\"}, ], \"answer\": \"John Doe\" }"
                     },
-                    "dataset_schema": {
+                    "source": {
-                        "type": "object",
+                        "$ref": "#/components/schemas/DataSource",
-                        "additionalProperties": {
+                        "description": "The data source of the dataset. Ensure that the data source schema is compatible with the purpose of the dataset. Examples: - { \"type\": \"uri\", \"uri\": \"https://mywebsite.com/mydata.jsonl\" } - { \"type\": \"uri\", \"uri\": \"lsfs://mydata.jsonl\" } - { \"type\": \"uri\", \"uri\": \"data:csv;base64,{base64_content}\" } - { \"type\": \"uri\", \"uri\": \"huggingface://llamastack/simpleqa?split=train\" } - { \"type\": \"rows\", \"rows\": [ { \"messages\": [ {\"role\": \"user\", \"content\": \"Hello, world!\"}, {\"role\": \"assistant\", \"content\": \"Hello, world!\"}, ] } ] }"
-                            "$ref": "#/components/schemas/ParamType"
-                        }
-                    },
-                    "url": {
-                        "$ref": "#/components/schemas/URL"
-                    },
-                    "provider_dataset_id": {
-                        "type": "string"
-                    },
-                    "provider_id": {
-                        "type": "string"
                     },
                     "metadata": {
                         "type": "object",
@@ -9371,14 +9454,18 @@
                                     "type": "object"
                                 }
                             ]
-                        }
+                        },
+                        "description": "The metadata for the dataset. - E.g. {\"description\": \"My dataset\"}"
+                    },
+                    "dataset_id": {
+                        "type": "string",
+                        "description": "The ID of the dataset. If not provided, an ID will be generated."
                     }
                 },
                 "additionalProperties": false,
                 "required": [
-                    "dataset_id",
+                    "purpose",
-                    "dataset_schema",
+                    "source"
-                    "url"
                 ],
                 "title": "RegisterDatasetRequest"
             },
@@ -10197,7 +10284,7 @@
             "x-displayName": "Llama Stack Evaluation API for running evaluations on model and agent candidates."
         },
         {
-            "name": "Files (Coming Soon)"
+            "name": "Files"
         },
         {
             "name": "Inference",
@@ -10258,7 +10345,7 @@
                 "DatasetIO",
                 "Datasets",
                 "Eval",
-                "Files (Coming Soon)",
+                "Files",
                 "Inference",
                 "Inspect",
                 "Models",

docs/_static/llama-stack-spec.yaml

@@ -10,56 +10,7 @@ info:
 servers:
   - url: http://any-hosted-llama-stack.com
 paths:
-  /v1/datasetio/rows:
+  /v1/datasetio/append-rows/{dataset_id}:
-    get:
-      responses:
-        '200':
-          description: OK
-          content:
-            application/json:
-              schema:
-                $ref: '#/components/schemas/PaginatedRowsResult'
-        '400':
-          $ref: '#/components/responses/BadRequest400'
-        '429':
-          $ref: >-
-            #/components/responses/TooManyRequests429
-        '500':
-          $ref: >-
-            #/components/responses/InternalServerError500
-        default:
-          $ref: '#/components/responses/DefaultError'
-      tags:
-        - DatasetIO
-      description: >-
-        Get a paginated list of rows from a dataset.
-      parameters:
-        - name: dataset_id
-          in: query
-          description: >-
-            The ID of the dataset to get the rows from.
-          required: true
-          schema:
-            type: string
-        - name: rows_in_page
-          in: query
-          description: The number of rows to get per page.
-          required: true
-          schema:
-            type: integer
-        - name: page_token
-          in: query
-          description: The token to get the next page of rows.
-          required: false
-          schema:
-            type: string
-        - name: filter_condition
-          in: query
-          description: >-
-            (Optional) A condition to filter the rows by.
-          required: false
-          schema:
-            type: string
     post:
       responses:
         '200':
@@ -77,7 +28,12 @@ paths:
       tags:
         - DatasetIO
       description: ''
-      parameters: []
+      parameters:
+        - name: dataset_id
+          in: path
+          required: true
+          schema:
+            type: string
       requestBody:
         content:
           application/json:
@@ -394,7 +350,7 @@ paths:
         default:
           $ref: '#/components/responses/DefaultError'
       tags:
-        - Files (Coming Soon)
+        - Files
       description: List all buckets.
       parameters:
         - name: bucket
@@ -421,7 +377,7 @@ paths:
         default:
           $ref: '#/components/responses/DefaultError'
       tags:
-        - Files (Coming Soon)
+        - Files
       description: >-
         Create a new upload session for a file identified by a bucket and key.
       parameters: []
@@ -580,7 +536,7 @@ paths:
         default:
           $ref: '#/components/responses/DefaultError'
       tags:
-        - Files (Coming Soon)
+        - Files
       description: >-
         Get a file info identified by a bucket and key.
       parameters:
@@ -616,7 +572,7 @@ paths:
         default:
           $ref: '#/components/responses/DefaultError'
       tags:
-        - Files (Coming Soon)
+        - Files
       description: >-
         Delete a file identified by a bucket and key.
       parameters:
@@ -1268,7 +1224,7 @@ paths:
         default:
           $ref: '#/components/responses/DefaultError'
       tags:
-        - Files (Coming Soon)
+        - Files
       description: >-
         Returns information about an existsing upload session
       parameters:
@@ -1299,7 +1255,7 @@ paths:
         default:
           $ref: '#/components/responses/DefaultError'
       tags:
-        - Files (Coming Soon)
+        - Files
       description: >-
         Upload file content to an existing upload session. On the server, request
         body will have the raw bytes that are uploaded.
@@ -1501,6 +1457,50 @@ paths:
             schema:
               $ref: '#/components/schemas/InvokeToolRequest'
         required: true
+  /v1/datasetio/iterrows/{dataset_id}:
+    get:
+      responses:
+        '200':
+          description: OK
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/IterrowsResponse'
+        '400':
+          $ref: '#/components/responses/BadRequest400'
+        '429':
+          $ref: >-
+            #/components/responses/TooManyRequests429
+        '500':
+          $ref: >-
+            #/components/responses/InternalServerError500
+        default:
+          $ref: '#/components/responses/DefaultError'
+      tags:
+        - DatasetIO
+      description: >-
+        Get a paginated list of rows from a dataset. Uses cursor-based pagination.
+      parameters:
+        - name: dataset_id
+          in: path
+          description: >-
+            The ID of the dataset to get the rows from.
+          required: true
+          schema:
+            type: string
+        - name: start_index
+          in: query
+          description: >-
+            Index into dataset for the first row to get. Get all rows if None.
+          required: false
+          schema:
+            type: integer
+        - name: limit
+          in: query
+          description: The number of rows to get.
+          required: false
+          schema:
+            type: integer
   /v1/eval/benchmarks/{benchmark_id}/jobs/{job_id}:
     get:
       responses:
@@ -1710,6 +1710,10 @@ paths:
       responses:
         '200':
           description: OK
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/Dataset'
         '400':
           $ref: '#/components/responses/BadRequest400'
         '429':
@@ -1722,7 +1726,7 @@ paths:
           $ref: '#/components/responses/DefaultError'
       tags:
         - Datasets
-      description: ''
+      description: Register a new dataset.
       parameters: []
       requestBody:
         content:
@@ -1750,7 +1754,7 @@ paths:
         default:
           $ref: '#/components/responses/DefaultError'
       tags:
-        - Files (Coming Soon)
+        - Files
       description: List all files in a bucket.
       parameters:
         - name: bucket
@@ -2607,8 +2611,6 @@ components:
     AppendRowsRequest:
       type: object
       properties:
-        dataset_id:
-          type: string
         rows:
           type: array
           items:
@@ -2623,7 +2625,6 @@ components:
                 - type: object
       additionalProperties: false
       required:
-        - dataset_id
         - rows
       title: AppendRowsRequest
     CompletionMessage:
@@ -4726,6 +4727,148 @@ components:
         - scoring_functions
         - metadata
       title: Benchmark
+    DataSource:
+      oneOf:
+        - $ref: '#/components/schemas/URIDataSource'
+        - $ref: '#/components/schemas/RowsDataSource'
+      discriminator:
+        propertyName: type
+        mapping:
+          uri: '#/components/schemas/URIDataSource'
+          rows: '#/components/schemas/RowsDataSource'
+    Dataset:
+      type: object
+      properties:
+        identifier:
+          type: string
+        provider_resource_id:
+          type: string
+        provider_id:
+          type: string
+        type:
+          type: string
+          const: dataset
+          default: dataset
+        purpose:
+          type: string
+          enum:
+            - post-training/messages
+            - eval/question-answer
+            - eval/messages-answer
+          title: DatasetPurpose
+          description: >-
+            Purpose of the dataset. Each purpose has a required input data schema.
+        source:
+          $ref: '#/components/schemas/DataSource'
+        metadata:
+          type: object
+          additionalProperties:
+            oneOf:
+              - type: 'null'
+              - type: boolean
+              - type: number
+              - type: string
+              - type: array
+              - type: object
+      additionalProperties: false
+      required:
+        - identifier
+        - provider_resource_id
+        - provider_id
+        - type
+        - purpose
+        - source
+        - metadata
+      title: Dataset
+    RowsDataSource:
+      type: object
+      properties:
+        type:
+          type: string
+          const: rows
+          default: rows
+        rows:
+          type: array
+          items:
+            type: object
+            additionalProperties:
+              oneOf:
+                - type: 'null'
+                - type: boolean
+                - type: number
+                - type: string
+                - type: array
+                - type: object
+          description: >-
+            The dataset is stored in rows. E.g. - [ {"messages": [{"role": "user",
+            "content": "Hello, world!"}, {"role": "assistant", "content": "Hello,
+            world!"}]} ]
+      additionalProperties: false
+      required:
+        - type
+        - rows
+      title: RowsDataSource
+      description: A dataset stored in rows.
+    URIDataSource:
+      type: object
+      properties:
+        type:
+          type: string
+          const: uri
+          default: uri
+        uri:
+          type: string
+          description: >-
+            The dataset can be obtained from a URI. E.g. - "https://mywebsite.com/mydata.jsonl"
+            - "lsfs://mydata.jsonl" - "data:csv;base64,{base64_content}"
+      additionalProperties: false
+      required:
+        - type
+        - uri
+      title: URIDataSource
+      description: >-
+        A dataset that can be obtained from a URI.
+    Model:
+      type: object
+      properties:
+        identifier:
+          type: string
+        provider_resource_id:
+          type: string
+        provider_id:
+          type: string
+        type:
+          type: string
+          const: model
+          default: model
+        metadata:
+          type: object
+          additionalProperties:
+            oneOf:
+              - type: 'null'
+              - type: boolean
+              - type: number
+              - type: string
+              - type: array
+              - type: object
+        model_type:
+          $ref: '#/components/schemas/ModelType'
+          default: llm
+      additionalProperties: false
+      required:
+        - identifier
+        - provider_resource_id
+        - provider_id
+        - type
+        - metadata
+        - model_type
+      title: Model
+    ModelType:
+      type: string
+      enum:
+        - llm
+        - embedding
+      title: ModelType
     AgentTurnInputType:
       type: object
       properties:
@@ -4781,45 +4924,6 @@ components:
       required:
         - type
       title: CompletionInputType
-    Dataset:
-      type: object
-      properties:
-        identifier:
-          type: string
-        provider_resource_id:
-          type: string
-        provider_id:
-          type: string
-        type:
-          type: string
-          const: dataset
-          default: dataset
-        dataset_schema:
-          type: object
-          additionalProperties:
-            $ref: '#/components/schemas/ParamType'
-        url:
-          $ref: '#/components/schemas/URL'
-        metadata:
-          type: object
-          additionalProperties:
-            oneOf:
-              - type: 'null'
-              - type: boolean
-              - type: number
-              - type: string
-              - type: array
-              - type: object
-      additionalProperties: false
-      required:
-        - identifier
-        - provider_resource_id
-        - provider_id
-        - type
-        - dataset_schema
-        - url
-        - metadata
-      title: Dataset
     JsonType:
       type: object
       properties:
@@ -4878,97 +4982,6 @@ components:
           chat_completion_input: '#/components/schemas/ChatCompletionInputType'
           completion_input: '#/components/schemas/CompletionInputType'
           agent_turn_input: '#/components/schemas/AgentTurnInputType'
-    StringType:
-      type: object
-      properties:
-        type:
-          type: string
-          const: string
-          default: string
-      additionalProperties: false
-      required:
-        - type
-      title: StringType
-    UnionType:
-      type: object
-      properties:
-        type:
-          type: string
-          const: union
-          default: union
-      additionalProperties: false
-      required:
-        - type
-      title: UnionType
-    Model:
-      type: object
-      properties:
-        identifier:
-          type: string
-        provider_resource_id:
-          type: string
-        provider_id:
-          type: string
-        type:
-          type: string
-          const: model
-          default: model
-        metadata:
-          type: object
-          additionalProperties:
-            oneOf:
-              - type: 'null'
-              - type: boolean
-              - type: number
-              - type: string
-              - type: array
-              - type: object
-        model_type:
-          $ref: '#/components/schemas/ModelType'
-          default: llm
-      additionalProperties: false
-      required:
-        - identifier
-        - provider_resource_id
-        - provider_id
-        - type
-        - metadata
-        - model_type
-      title: Model
-    ModelType:
-      type: string
-      enum:
-        - llm
-        - embedding
-      title: ModelType
-    PaginatedRowsResult:
-      type: object
-      properties:
-        rows:
-          type: array
-          items:
-            type: object
-            additionalProperties:
-              oneOf:
-                - type: 'null'
-                - type: boolean
-                - type: number
-                - type: string
-                - type: array
-                - type: object
-          description: The rows in the current page.
-        total_count:
-          type: integer
-          description: The total number of rows in the dataset.
-        next_page_token:
-          type: string
-          description: The token to get the next page of rows.
-      additionalProperties: false
-      required:
-        - rows
-        - total_count
-      title: PaginatedRowsResult
-      description: A paginated list of rows from a dataset.
     ScoringFn:
       type: object
       properties:
@@ -5007,6 +5020,28 @@ components:
         - metadata
         - return_type
       title: ScoringFn
+    StringType:
+      type: object
+      properties:
+        type:
+          type: string
+          const: string
+          default: string
+      additionalProperties: false
+      required:
+        - type
+      title: StringType
+    UnionType:
+      type: object
+      properties:
+        type:
+          type: string
+          const: union
+          default: union
+      additionalProperties: false
+      required:
+        - type
+      title: UnionType
     Shield:
       type: object
       properties:
@@ -5506,6 +5541,32 @@ components:
       required:
         - content
       title: ToolInvocationResult
+    IterrowsResponse:
+      type: object
+      properties:
+        data:
+          type: array
+          items:
+            type: object
+            additionalProperties:
+              oneOf:
+                - type: 'null'
+                - type: boolean
+                - type: number
+                - type: string
+                - type: array
+                - type: object
+          description: The rows in the current page.
+        next_start_index:
+          type: integer
+          description: >-
+            Index into dataset for the first row in the next page. None if there are
+            no more rows.
+      additionalProperties: false
+      required:
+        - data
+      title: IterrowsResponse
+      description: A paginated list of rows from a dataset.
     ListAgentSessionsResponse:
       type: object
       properties:
@@ -6313,18 +6374,35 @@ components:
     RegisterDatasetRequest:
       type: object
       properties:
-        dataset_id:
+        purpose:
-          type: string
-        dataset_schema:
-          type: object
-          additionalProperties:
-            $ref: '#/components/schemas/ParamType'
-        url:
-          $ref: '#/components/schemas/URL'
-        provider_dataset_id:
-          type: string
-        provider_id:
           type: string
+          enum:
+            - post-training/messages
+            - eval/question-answer
+            - eval/messages-answer
+          description: >-
+            The purpose of the dataset. One of - "post-training/messages": The dataset
+            contains a messages column with list of messages for post-training. {
+            "messages": [ {"role": "user", "content": "Hello, world!"}, {"role": "assistant",
+            "content": "Hello, world!"}, ] } - "eval/question-answer": The dataset
+            contains a question column and an answer column for evaluation. { "question":
+            "What is the capital of France?", "answer": "Paris" } - "eval/messages-answer":
+            The dataset contains a messages column with list of messages and an answer
+            column for evaluation. { "messages": [ {"role": "user", "content": "Hello,
+            my name is John Doe."}, {"role": "assistant", "content": "Hello, John
+            Doe. How can I help you today?"}, {"role": "user", "content": "What's
+            my name?"}, ], "answer": "John Doe" }
+        source:
+          $ref: '#/components/schemas/DataSource'
+          description: >-
+            The data source of the dataset. Ensure that the data source schema is
+            compatible with the purpose of the dataset. Examples: - { "type": "uri",
+            "uri": "https://mywebsite.com/mydata.jsonl" } - { "type": "uri", "uri":
+            "lsfs://mydata.jsonl" } - { "type": "uri", "uri": "data:csv;base64,{base64_content}"
+            } - { "type": "uri", "uri": "huggingface://llamastack/simpleqa?split=train"
+            } - { "type": "rows", "rows": [ { "messages": [ {"role": "user", "content":
+            "Hello, world!"}, {"role": "assistant", "content": "Hello, world!"}, ]
+            } ] }
         metadata:
           type: object
           additionalProperties:
@@ -6335,11 +6413,16 @@ components:
               - type: string
               - type: array
               - type: object
+          description: >-
+            The metadata for the dataset. - E.g. {"description": "My dataset"}
+        dataset_id:
+          type: string
+          description: >-
+            The ID of the dataset. If not provided, an ID will be generated.
       additionalProperties: false
       required:
-        - dataset_id
+        - purpose
-        - dataset_schema
+        - source
-        - url
       title: RegisterDatasetRequest
     RegisterModelRequest:
       type: object
@@ -6855,7 +6938,7 @@ tags:
   - name: Eval
     x-displayName: >-
       Llama Stack Evaluation API for running evaluations on model and agent candidates.
-  - name: Files (Coming Soon)
+  - name: Files
   - name: Inference
     description: >-
       This API provides the raw interface to the underlying models. Two kinds of models
@@ -6893,7 +6976,7 @@ x-tagGroups:
       - DatasetIO
       - Datasets
       - Eval
-      - Files (Coming Soon)
+      - Files
       - Inference
       - Inspect
       - Models

docs/notebooks/Llama_Stack_Benchmark_Evals.ipynb

@@ -84,16 +84,14 @@
           "name": "stdout",
           "output_type": "stream",
           "text": [
-            "Not in Google Colab environment\n",
+            "Not in Google Colab environment\n"
-            "\u001b[33mWarning: `bwrap` is not available. Code interpreter tool will not work correctly.\u001b[0m\n"
           ]
         },
         {
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "/opt/anaconda3/envs/master/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+            "Warning: `bwrap` is not available. Code interpreter tool will not work correctly.\n"
-            "  from .autonotebook import tqdm as notebook_tqdm\n"
           ]
         },
         {
@@ -117,76 +115,146 @@
               "- datasetio\n",
               "- eval\n",
               "- inference\n",
-              "- memory\n",
               "- safety\n",
               "- scoring\n",
               "- telemetry\n",
               "- tool_runtime\n",
-              "datasets: <span style=\"font-weight: bold\">[]</span>\n",
+              "- vector_io\n",
-              "container_image: null\n",
               "benchmarks: <span style=\"font-weight: bold\">[]</span>\n",
+              "container_image: null\n",
+              "datasets: <span style=\"font-weight: bold\">[]</span>\n",
               "image_name: together\n",
-              "memory_banks: <span style=\"font-weight: bold\">[]</span>\n",
+              "logging: null\n",
               "metadata_store:\n",
               "  db_path: <span style=\"color: #800080; text-decoration-color: #800080\">/Users/xiyan/.llama/distributions/together/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">registry.db</span>\n",
               "  namespace: null\n",
               "  type: sqlite\n",
               "models:\n",
               "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
+              "  model_id: meta-llama/Meta-Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-8B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Meta-Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-8B-Instruct-Turbo\n",
+              "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
               "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-8B-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Meta-Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-8B-Instruct-Turbo\n",
               "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
+              "  model_id: meta-llama/Meta-Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-70B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Meta-Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-70B-Instruct-Turbo\n",
+              "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
               "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-70B-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Meta-Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-70B-Instruct-Turbo\n",
               "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
+              "  model_id: meta-llama/Meta-Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-405B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Meta-Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-405B-Instruct-Turbo\n",
+              "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
               "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-405B-Instruct-FP8\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Meta-Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.1</span>-405B-Instruct-Turbo\n",
               "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
+              "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-3B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-3B-Instruct-Turbo\n",
+              "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
               "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-3B-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-3B-Instruct-Turbo\n",
               "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
+              "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-11B-Vision-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-11B-Vision-Instruct-Turbo\n",
+              "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
               "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-11B-Vision-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-11B-Vision-Instruct-Turbo\n",
               "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
+              "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-90B-Vision-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-90B-Vision-Instruct-Turbo\n",
+              "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
               "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-90B-Vision-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.2</span>-90B-Vision-Instruct-Turbo\n",
               "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
+              "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.3</span>-70B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.3</span>-70B-Instruct-Turbo\n",
+              "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
               "  model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.3</span>-70B-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3.3</span>-70B-Instruct-Turbo\n",
               "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
+              "  model_id: meta-llama/Meta-Llama-Guard-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>-8B\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Meta-Llama-Guard-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>-8B\n",
+              "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
               "  model_id: meta-llama/Llama-Guard-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>-8B\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Meta-Llama-Guard-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>-8B\n",
               "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
+              "  model_id: meta-llama/Llama-Guard-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>-11B-Vision-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-Guard-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>-11B-Vision-Turbo\n",
+              "- metadata: <span style=\"font-weight: bold\">{}</span>\n",
               "  model_id: meta-llama/Llama-Guard-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>-11B-Vision\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-Guard-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">3</span>-11B-Vision-Turbo\n",
               "- metadata:\n",
+              "    context_length: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8192</span>\n",
+              "    embedding_dimension: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">768</span>\n",
+              "  model_id: togethercomputer/m2-bert-80M-8k-retrieval\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - embedding\n",
+              "  provider_id: together\n",
+              "  provider_model_id: togethercomputer/m2-bert-80M-8k-retrieval\n",
+              "- metadata:\n",
+              "    context_length: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">32768</span>\n",
+              "    embedding_dimension: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">768</span>\n",
+              "  model_id: togethercomputer/m2-bert-80M-32k-retrieval\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - embedding\n",
+              "  provider_id: together\n",
+              "  provider_model_id: togethercomputer/m2-bert-80M-32k-retrieval\n",
+              "- metadata:\n",
               "    embedding_dimension: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">384</span>\n",
               "  model_id: all-MiniLM-L6-v2\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
@@ -203,14 +271,26 @@
               "    provider_id: meta-reference\n",
               "    provider_type: inline::meta-reference\n",
               "  datasetio:\n",
-              "  - config: <span style=\"font-weight: bold\">{}</span>\n",
+              "  - config:\n",
+              "      kvstore:\n",
+              "        db_path: <span style=\"color: #800080; text-decoration-color: #800080\">/Users/xiyan/.llama/distributions/together/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">huggingface_datasetio.db</span>\n",
+              "        namespace: null\n",
+              "        type: sqlite\n",
               "    provider_id: huggingface\n",
               "    provider_type: remote::huggingface\n",
-              "  - config: <span style=\"font-weight: bold\">{}</span>\n",
+              "  - config:\n",
+              "      kvstore:\n",
+              "        db_path: <span style=\"color: #800080; text-decoration-color: #800080\">/Users/xiyan/.llama/distributions/together/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">localfs_datasetio.db</span>\n",
+              "        namespace: null\n",
+              "        type: sqlite\n",
               "    provider_id: localfs\n",
               "    provider_type: inline::localfs\n",
               "  eval:\n",
-              "  - config: <span style=\"font-weight: bold\">{}</span>\n",
+              "  - config:\n",
+              "      kvstore:\n",
+              "        db_path: <span style=\"color: #800080; text-decoration-color: #800080\">/Users/xiyan/.llama/distributions/together/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">meta_reference_eval.db</span>\n",
+              "        namespace: null\n",
+              "        type: sqlite\n",
               "    provider_id: meta-reference\n",
               "    provider_type: inline::meta-reference\n",
               "  inference:\n",
@@ -222,16 +302,9 @@
               "  - config: <span style=\"font-weight: bold\">{}</span>\n",
               "    provider_id: sentence-transformers\n",
               "    provider_type: inline::sentence-transformers\n",
-              "  memory:\n",
-              "  - config:\n",
-              "      kvstore:\n",
-              "        db_path: <span style=\"color: #800080; text-decoration-color: #800080\">/Users/xiyan/.llama/distributions/together/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">faiss_store.db</span>\n",
-              "        namespace: null\n",
-              "        type: sqlite\n",
-              "    provider_id: faiss\n",
-              "    provider_type: inlin<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\">e::fa</span>iss\n",
               "  safety:\n",
-              "  - config: <span style=\"font-weight: bold\">{}</span>\n",
+              "  - config:\n",
+              "      excluded_categories: <span style=\"font-weight: bold\">[]</span>\n",
               "    provider_id: llama-guard\n",
               "    provider_type: inline::llama-guard\n",
               "  scoring:\n",
@@ -269,7 +342,26 @@
               "  - config: <span style=\"font-weight: bold\">{}</span>\n",
               "    provider_id: rag-runtime\n",
               "    provider_type: inline::rag-runtime\n",
+              "  - config: <span style=\"font-weight: bold\">{}</span>\n",
+              "    provider_id: model-context-protocol\n",
+              "    provider_type: remote::model-context-protocol\n",
+              "  - config:\n",
+              "      api_key: <span style=\"color: #008000; text-decoration-color: #008000\">'********'</span>\n",
+              "    provider_id: wolfram-alpha\n",
+              "    provider_type: remote::wolfram-alpha\n",
+              "  vector_io:\n",
+              "  - config:\n",
+              "      kvstore:\n",
+              "        db_path: <span style=\"color: #800080; text-decoration-color: #800080\">/Users/xiyan/.llama/distributions/together/</span><span style=\"color: #ff00ff; text-decoration-color: #ff00ff\">faiss_store.db</span>\n",
+              "        namespace: null\n",
+              "        type: sqlite\n",
+              "    provider_id: faiss\n",
+              "    provider_type: inlin<span style=\"color: #00ff00; text-decoration-color: #00ff00; font-weight: bold\">e::fa</span>iss\n",
               "scoring_fns: <span style=\"font-weight: bold\">[]</span>\n",
+              "server:\n",
+              "  port: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8321</span>\n",
+              "  tls_certfile: null\n",
+              "  tls_keyfile: null\n",
               "shields:\n",
               "- params: null\n",
               "  provider_id: null\n",
@@ -288,6 +380,11 @@
               "  mcp_endpoint: null\n",
               "  provider_id: code-interpreter\n",
               "  toolgroup_id: builtin::code_interpreter\n",
+              "- args: null\n",
+              "  mcp_endpoint: null\n",
+              "  provider_id: wolfram-alpha\n",
+              "  toolgroup_id: builtin::wolfram_alpha\n",
+              "vector_dbs: <span style=\"font-weight: bold\">[]</span>\n",
               "version: <span style=\"color: #008000; text-decoration-color: #008000\">'2'</span>\n",
               "\n",
               "</pre>\n"
@@ -298,76 +395,146 @@
               "- datasetio\n",
               "- eval\n",
               "- inference\n",
-              "- memory\n",
               "- safety\n",
               "- scoring\n",
               "- telemetry\n",
               "- tool_runtime\n",
-              "datasets: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+              "- vector_io\n",
-              "container_image: null\n",
               "benchmarks: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+              "container_image: null\n",
+              "datasets: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
               "image_name: together\n",
-              "memory_banks: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+              "logging: null\n",
               "metadata_store:\n",
               "  db_path: \u001b[35m/Users/xiyan/.llama/distributions/together/\u001b[0m\u001b[95mregistry.db\u001b[0m\n",
               "  namespace: null\n",
               "  type: sqlite\n",
               "models:\n",
               "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  model_id: meta-llama/Meta-Llama-\u001b[1;36m3.1\u001b[0m-8B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Meta-Llama-\u001b[1;36m3.1\u001b[0m-8B-Instruct-Turbo\n",
+              "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "  model_id: meta-llama/Llama-\u001b[1;36m3.1\u001b[0m-8B-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Meta-Llama-\u001b[1;36m3.1\u001b[0m-8B-Instruct-Turbo\n",
               "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  model_id: meta-llama/Meta-Llama-\u001b[1;36m3.1\u001b[0m-70B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Meta-Llama-\u001b[1;36m3.1\u001b[0m-70B-Instruct-Turbo\n",
+              "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "  model_id: meta-llama/Llama-\u001b[1;36m3.1\u001b[0m-70B-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Meta-Llama-\u001b[1;36m3.1\u001b[0m-70B-Instruct-Turbo\n",
               "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  model_id: meta-llama/Meta-Llama-\u001b[1;36m3.1\u001b[0m-405B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Meta-Llama-\u001b[1;36m3.1\u001b[0m-405B-Instruct-Turbo\n",
+              "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "  model_id: meta-llama/Llama-\u001b[1;36m3.1\u001b[0m-405B-Instruct-FP8\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Meta-Llama-\u001b[1;36m3.1\u001b[0m-405B-Instruct-Turbo\n",
               "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-3B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-3B-Instruct-Turbo\n",
+              "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "  model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-3B-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-3B-Instruct-Turbo\n",
               "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-11B-Vision-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-11B-Vision-Instruct-Turbo\n",
+              "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "  model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-11B-Vision-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-11B-Vision-Instruct-Turbo\n",
               "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-90B-Vision-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-90B-Vision-Instruct-Turbo\n",
+              "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "  model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-90B-Vision-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-\u001b[1;36m3.2\u001b[0m-90B-Vision-Instruct-Turbo\n",
               "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  model_id: meta-llama/Llama-\u001b[1;36m3.3\u001b[0m-70B-Instruct-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-\u001b[1;36m3.3\u001b[0m-70B-Instruct-Turbo\n",
+              "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "  model_id: meta-llama/Llama-\u001b[1;36m3.3\u001b[0m-70B-Instruct\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-\u001b[1;36m3.3\u001b[0m-70B-Instruct-Turbo\n",
               "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  model_id: meta-llama/Meta-Llama-Guard-\u001b[1;36m3\u001b[0m-8B\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Meta-Llama-Guard-\u001b[1;36m3\u001b[0m-8B\n",
+              "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "  model_id: meta-llama/Llama-Guard-\u001b[1;36m3\u001b[0m-8B\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Meta-Llama-Guard-\u001b[1;36m3\u001b[0m-8B\n",
               "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  model_id: meta-llama/Llama-Guard-\u001b[1;36m3\u001b[0m-11B-Vision-Turbo\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - llm\n",
+              "  provider_id: together\n",
+              "  provider_model_id: meta-llama/Llama-Guard-\u001b[1;36m3\u001b[0m-11B-Vision-Turbo\n",
+              "- metadata: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "  model_id: meta-llama/Llama-Guard-\u001b[1;36m3\u001b[0m-11B-Vision\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
               "  - llm\n",
               "  provider_id: together\n",
               "  provider_model_id: meta-llama/Llama-Guard-\u001b[1;36m3\u001b[0m-11B-Vision-Turbo\n",
               "- metadata:\n",
+              "    context_length: \u001b[1;36m8192\u001b[0m\n",
+              "    embedding_dimension: \u001b[1;36m768\u001b[0m\n",
+              "  model_id: togethercomputer/m2-bert-80M-8k-retrieval\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - embedding\n",
+              "  provider_id: together\n",
+              "  provider_model_id: togethercomputer/m2-bert-80M-8k-retrieval\n",
+              "- metadata:\n",
+              "    context_length: \u001b[1;36m32768\u001b[0m\n",
+              "    embedding_dimension: \u001b[1;36m768\u001b[0m\n",
+              "  model_id: togethercomputer/m2-bert-80M-32k-retrieval\n",
+              "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
+              "  - embedding\n",
+              "  provider_id: together\n",
+              "  provider_model_id: togethercomputer/m2-bert-80M-32k-retrieval\n",
+              "- metadata:\n",
               "    embedding_dimension: \u001b[1;36m384\u001b[0m\n",
               "  model_id: all-MiniLM-L6-v2\n",
               "  model_type: !!python/object/apply:llama_stack.apis.models.models.ModelType\n",
@@ -384,14 +551,26 @@
               "    provider_id: meta-reference\n",
               "    provider_type: inline::meta-reference\n",
               "  datasetio:\n",
-              "  - config: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  - config:\n",
+              "      kvstore:\n",
+              "        db_path: \u001b[35m/Users/xiyan/.llama/distributions/together/\u001b[0m\u001b[95mhuggingface_datasetio.db\u001b[0m\n",
+              "        namespace: null\n",
+              "        type: sqlite\n",
               "    provider_id: huggingface\n",
               "    provider_type: remote::huggingface\n",
-              "  - config: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  - config:\n",
+              "      kvstore:\n",
+              "        db_path: \u001b[35m/Users/xiyan/.llama/distributions/together/\u001b[0m\u001b[95mlocalfs_datasetio.db\u001b[0m\n",
+              "        namespace: null\n",
+              "        type: sqlite\n",
               "    provider_id: localfs\n",
               "    provider_type: inline::localfs\n",
               "  eval:\n",
-              "  - config: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  - config:\n",
+              "      kvstore:\n",
+              "        db_path: \u001b[35m/Users/xiyan/.llama/distributions/together/\u001b[0m\u001b[95mmeta_reference_eval.db\u001b[0m\n",
+              "        namespace: null\n",
+              "        type: sqlite\n",
               "    provider_id: meta-reference\n",
               "    provider_type: inline::meta-reference\n",
               "  inference:\n",
@@ -403,16 +582,9 @@
               "  - config: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "    provider_id: sentence-transformers\n",
               "    provider_type: inline::sentence-transformers\n",
-              "  memory:\n",
-              "  - config:\n",
-              "      kvstore:\n",
-              "        db_path: \u001b[35m/Users/xiyan/.llama/distributions/together/\u001b[0m\u001b[95mfaiss_store.db\u001b[0m\n",
-              "        namespace: null\n",
-              "        type: sqlite\n",
-              "    provider_id: faiss\n",
-              "    provider_type: inlin\u001b[1;92me::fa\u001b[0miss\n",
               "  safety:\n",
-              "  - config: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "  - config:\n",
+              "      excluded_categories: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
               "    provider_id: llama-guard\n",
               "    provider_type: inline::llama-guard\n",
               "  scoring:\n",
@@ -450,7 +622,26 @@
               "  - config: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
               "    provider_id: rag-runtime\n",
               "    provider_type: inline::rag-runtime\n",
+              "  - config: \u001b[1m{\u001b[0m\u001b[1m}\u001b[0m\n",
+              "    provider_id: model-context-protocol\n",
+              "    provider_type: remote::model-context-protocol\n",
+              "  - config:\n",
+              "      api_key: \u001b[32m'********'\u001b[0m\n",
+              "    provider_id: wolfram-alpha\n",
+              "    provider_type: remote::wolfram-alpha\n",
+              "  vector_io:\n",
+              "  - config:\n",
+              "      kvstore:\n",
+              "        db_path: \u001b[35m/Users/xiyan/.llama/distributions/together/\u001b[0m\u001b[95mfaiss_store.db\u001b[0m\n",
+              "        namespace: null\n",
+              "        type: sqlite\n",
+              "    provider_id: faiss\n",
+              "    provider_type: inlin\u001b[1;92me::fa\u001b[0miss\n",
               "scoring_fns: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
+              "server:\n",
+              "  port: \u001b[1;36m8321\u001b[0m\n",
+              "  tls_certfile: null\n",
+              "  tls_keyfile: null\n",
               "shields:\n",
               "- params: null\n",
               "  provider_id: null\n",
@@ -469,6 +660,11 @@
               "  mcp_endpoint: null\n",
               "  provider_id: code-interpreter\n",
               "  toolgroup_id: builtin::code_interpreter\n",
+              "- args: null\n",
+              "  mcp_endpoint: null\n",
+              "  provider_id: wolfram-alpha\n",
+              "  toolgroup_id: builtin::wolfram_alpha\n",
+              "vector_dbs: \u001b[1m[\u001b[0m\u001b[1m]\u001b[0m\n",
               "version: \u001b[32m'2'\u001b[0m\n",
               "\n"
             ]
@@ -532,7 +728,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 3,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -643,17 +839,7 @@
         "id": "DJkmoG2kq1_P",
         "outputId": "8493ee59-c6ff-4bb6-d787-f295944db1cf"
       },
-      "outputs": [
+      "outputs": [],
-        {
-          "name": "stderr",
-          "output_type": "stream",
-          "text": [
-            "Generating dev split: 100%|██████████| 5/5 [00:00<00:00, 139.81 examples/s]\n",
-            "Generating validation split: 100%|██████████| 30/30 [00:00<00:00, 258.29 examples/s]\n",
-            "Generating test split: 100%|██████████| 287/287 [00:01<00:00, 197.69 examples/s]\n"
-          ]
-        }
-      ],
       "source": [
         "import datasets\n",
         "\n",
@@ -676,7 +862,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 7,
+      "execution_count": 4,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -691,7 +877,7 @@
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "100%|██████████| 5/5 [00:42<00:00,  8.60s/it]\n"
+            "100%|██████████| 5/5 [00:33<00:00,  6.71s/it]\n"
           ]
         },
         {
@@ -699,16 +885,18 @@
             "text/html": [
               "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">EvaluateResponse</span><span style=\"font-weight: bold\">(</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #808000; text-decoration-color: #808000\">generations</span>=<span style=\"font-weight: bold\">[</span>\n",
-              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'Answer: D'</span><span style=\"font-weight: bold\">}</span>,\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span>\n",
-              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'The image shows a sunflower leaf with small, dark spots and white powdery patches. The dark spots are likely caused by a fungal pathogen, such as rust or septoria leaf spot, while the white powdery patches are likely caused by a fungal pathogen, such as powdery mildew.\\n\\nSince there are two distinct types of lesions on the leaf, it is likely that there are two different pathogens infecting the leaf.\\n\\n**Answer:** B) Two pathogens'</span>\n",
+              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'**Potato Pests**\\n\\nThe two insects depicted are:\\n\\n* **Colorado Potato Beetle (Leptinotarsa decemlineata)**: Characterized by black and yellow stripes, this beetle is a significant pest of potatoes. It feeds on the leaves and can cause substantial damage to the crop.\\n* **False Potato Beetle (Leptinotarsa juncta)**: Also known as the false Colorado beetle, this species has similar coloring but is not as harmful to potatoes as the Colorado potato beetle.'</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>,\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span>\n",
-              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">\"The question requires the identification of the reason behind the massive gum production on the trunks of grapefruit trees in Cyprus, despite appearing healthy from a distance. The correct answer can be deduced by analyzing the symptoms and considering the possible causes.\\n\\nTo determine the correct answer, let's evaluate each option:\\n\\nA) Don't know or not sure: This option is incorrect because it does not provide a specific reason for the gum production.\\n\\nB) Physiological stress: This option is also incorrect because it is too broad and does not specifically explain the gum production.\\n\\nC) Bacterial disease: This option is incorrect because bacterial diseases typically cause different symptoms such as leaf spots, blights, or wilting.\\n\\nD) Harvesting damage when cutting with knives: This option is incorrect because harvesting damage would likely cause wounds or scars on the tree, but it would not lead to massive gum production.\\n\\nE) Fungal gummosis: This option is the most likely cause of the gum production. Fungal gummosis is a common disease in citrus trees, including grapefruit, that causes the production of gum or sap on the trunks and branches. The disease is typically caused by fungi such as Phytophthora or Diplodia, which infect the tree through wounds or natural openings. The gum production is a defense mechanism by the tree to try to seal off the infection and prevent further damage.\\n\\nTherefore, the correct answer is:\\n\\nAnswer: E\"</span>\n",
+              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">\"The image shows a sunflower leaf with a powdery mildew, which is a fungal disease caused by various species of fungi. The white powdery coating on the leaves is a characteristic symptom of this disease. The leaf also has some black spots, which could be indicative of a secondary infection or another type of disease. However, without more information or a closer examination, it's difficult to determine the exact cause of the black spots.\\n\\nBased on the image alone, we can see at least two types of symptoms: the powdery mildew and the black spots. This suggests that there may be more than one pathogen involved, but it's also possible that the black spots are a result of the same fungal infection causing the powdery mildew.\\n\\nAnswer: B) Two pathogens\"</span>\n",
+              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>,\n",
+              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span>\n",
+              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'The symptoms observed, characterized by the massive gum production on the trunks of the grapefruit trees in Cyprus, suggest a physiological or pathological response. Given the absence of visible signs of damage or pests from a higher point on a hillside, and considering the specific nature of the symptom (gum production), we can infer that the cause is more likely related to an internal process within the tree rather than external damage from harvesting. While physiological stress (B) could lead to such symptoms, the primary reason for gum production in trees, especially in citrus species, is typically linked to disease. Among the options provided, fungal gummosis (E) is a condition known to cause gumming in citrus trees, which aligns with the observed symptoms. Therefore, without direct evidence of external damage (harvesting) or confirmation of physiological stress being the primary cause, the most appropriate answer based on the information given is:\\n\\nAnswer: E'</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>,\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'Answer: D'</span><span style=\"font-weight: bold\">}</span>,\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span>\n",
-              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'**Causes of Splitting Petioles in Rhubarb**\\n\\nThe following factors can cause the petioles of rhubarb to split:\\n\\n* **Physiological Problems**: Issues such as water stress, nutrient deficiencies, or extreme temperatures can lead to splitting.\\n* **Phytoplasma Infection**: A bacterial infection caused by phytoplasma can lead to splitting of the petioles.\\n* **Animal Damage**: Pests like slugs, snails, or rodents can damage the plant and cause splitting.\\n* **Bacterial Infection**: Bacterial infections can also cause splitting.\\n\\nAs a result, the correct answer is:\\n\\n*Answer*: A) Physiological problems'</span>\n",
+              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">\"**Analysis of the Image**\\n\\nThe image provided shows a rhubarb plant with split petioles. To determine the cause of this issue, we need to consider various factors that could lead to such damage.\\n\\n**Possible Causes of Petiole Splitting**\\n\\n*   **Physiological Problems**: Rhubarb plants can experience physiological stress due to environmental factors like extreme temperatures, waterlogging, or nutrient deficiencies. This stress can cause the petioles to split.\\n*   **Phytoplasma Infection**: Phytoplasma is a type of bacteria that can infect plants, including rhubarb. It can cause symptoms such as yellowing leaves, stunted growth, and splitting of petioles.\\n*   **Animal Damage**: Animals like rabbits, deer, or insects can damage rhubarb plants by eating the leaves or stems, which can lead to splitting of the petioles.\\n*   **Bacteria**: Bacterial infections can also cause damage to rhubarb plants, including splitting of the petioles.\\n\\n**Conclusion**\\n\\nBased on the analysis, it is clear that all the options listed (A) Physiological problems, B) Phytoplasma infection, D) Animal damage, and E) Bacteria) could potentially cause the petioles of the rhubarb plant to split. Therefore, there is no single option that would not be a cause for the petioles splitting.\\n\\n**Answer**: C) I don't know and don't want to guess.\"</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"font-weight: bold\">]</span>,\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #808000; text-decoration-color: #808000\">scores</span>=<span style=\"font-weight: bold\">{</span>\n",
@@ -723,16 +911,18 @@
             "text/plain": [
               "\u001b[1;35mEvaluateResponse\u001b[0m\u001b[1m(\u001b[0m\n",
               "\u001b[2;32m│   \u001b[0m\u001b[33mgenerations\u001b[0m=\u001b[1m[\u001b[0m\n",
-              "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m'Answer: D'\u001b[0m\u001b[1m}\u001b[0m,\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\n",
-              "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m'The image shows a sunflower leaf with small, dark spots and white powdery patches. The dark spots are likely caused by a fungal pathogen, such as rust or septoria leaf spot, while the white powdery patches are likely caused by a fungal pathogen, such as powdery mildew.\\n\\nSince there are two distinct types of lesions on the leaf, it is likely that there are two different pathogens infecting the leaf.\\n\\n**Answer:** B\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Two pathogens'\u001b[0m\n",
+              "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m'**Potato Pests**\\n\\nThe two insects depicted are:\\n\\n* **Colorado Potato Beetle \u001b[0m\u001b[32m(\u001b[0m\u001b[32mLeptinotarsa decemlineata\u001b[0m\u001b[32m)\u001b[0m\u001b[32m**: Characterized by black and yellow stripes, this beetle is a significant pest of potatoes. It feeds on the leaves and can cause substantial damage to the crop.\\n* **False Potato Beetle \u001b[0m\u001b[32m(\u001b[0m\u001b[32mLeptinotarsa juncta\u001b[0m\u001b[32m)\u001b[0m\u001b[32m**: Also known as the false Colorado beetle, this species has similar coloring but is not as harmful to potatoes as the Colorado potato beetle.'\u001b[0m\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m,\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\n",
-              "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m\"The question requires the identification of the reason behind the massive gum production on the trunks of grapefruit trees in Cyprus, despite appearing healthy from a distance. The correct answer can be deduced by analyzing the symptoms and considering the possible causes.\\n\\nTo determine the correct answer, let's evaluate each option:\\n\\nA\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Don't know or not sure: This option is incorrect because it does not provide a specific reason for the gum production.\\n\\nB\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Physiological stress: This option is also incorrect because it is too broad and does not specifically explain the gum production.\\n\\nC\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Bacterial disease: This option is incorrect because bacterial diseases typically cause different symptoms such as leaf spots, blights, or wilting.\\n\\nD\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Harvesting damage when cutting with knives: This option is incorrect because harvesting damage would likely cause wounds or scars on the tree, but it would not lead to massive gum production.\\n\\nE\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Fungal gummosis: This option is the most likely cause of the gum production. Fungal gummosis is a common disease in citrus trees, including grapefruit, that causes the production of gum or sap on the trunks and branches. The disease is typically caused by fungi such as Phytophthora or Diplodia, which infect the tree through wounds or natural openings. The gum production is a defense mechanism by the tree to try to seal off the infection and prevent further damage.\\n\\nTherefore, the correct answer is:\\n\\nAnswer: E\"\u001b[0m\n",
+              "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m\"The image shows a sunflower leaf with a powdery mildew, which is a fungal disease caused by various species of fungi. The white powdery coating on the leaves is a characteristic symptom of this disease. The leaf also has some black spots, which could be indicative of a secondary infection or another type of disease. However, without more information or a closer examination, it's difficult to determine the exact cause of the black spots.\\n\\nBased on the image alone, we can see at least two types of symptoms: the powdery mildew and the black spots. This suggests that there may be more than one pathogen involved, but it's also possible that the black spots are a result of the same fungal infection causing the powdery mildew.\\n\\nAnswer: B\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Two pathogens\"\u001b[0m\n",
+              "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m,\n",
+              "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\n",
+              "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m'The symptoms observed, characterized by the massive gum production on the trunks of the grapefruit trees in Cyprus, suggest a physiological or pathological response. Given the absence of visible signs of damage or pests from a higher point on a hillside, and considering the specific nature of the symptom \u001b[0m\u001b[32m(\u001b[0m\u001b[32mgum production\u001b[0m\u001b[32m)\u001b[0m\u001b[32m, we can infer that the cause is more likely related to an internal process within the tree rather than external damage from harvesting. While physiological stress \u001b[0m\u001b[32m(\u001b[0m\u001b[32mB\u001b[0m\u001b[32m)\u001b[0m\u001b[32m could lead to such symptoms, the primary reason for gum production in trees, especially in citrus species, is typically linked to disease. Among the options provided, fungal gummosis \u001b[0m\u001b[32m(\u001b[0m\u001b[32mE\u001b[0m\u001b[32m)\u001b[0m\u001b[32m is a condition known to cause gumming in citrus trees, which aligns with the observed symptoms. Therefore, without direct evidence of external damage \u001b[0m\u001b[32m(\u001b[0m\u001b[32mharvesting\u001b[0m\u001b[32m)\u001b[0m\u001b[32m or confirmation of physiological stress being the primary cause, the most appropriate answer based on the information given is:\\n\\nAnswer: E'\u001b[0m\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m,\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m'Answer: D'\u001b[0m\u001b[1m}\u001b[0m,\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\n",
-              "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m'**Causes of Splitting Petioles in Rhubarb**\\n\\nThe following factors can cause the petioles of rhubarb to split:\\n\\n* **Physiological Problems**: Issues such as water stress, nutrient deficiencies, or extreme temperatures can lead to splitting.\\n* **Phytoplasma Infection**: A bacterial infection caused by phytoplasma can lead to splitting of the petioles.\\n* **Animal Damage**: Pests like slugs, snails, or rodents can damage the plant and cause splitting.\\n* **Bacterial Infection**: Bacterial infections can also cause splitting.\\n\\nAs a result, the correct answer is:\\n\\n*Answer*: A\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Physiological problems'\u001b[0m\n",
+              "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m\"**Analysis of the Image**\\n\\nThe image provided shows a rhubarb plant with split petioles. To determine the cause of this issue, we need to consider various factors that could lead to such damage.\\n\\n**Possible Causes of Petiole Splitting**\\n\\n*   **Physiological Problems**: Rhubarb plants can experience physiological stress due to environmental factors like extreme temperatures, waterlogging, or nutrient deficiencies. This stress can cause the petioles to split.\\n*   **Phytoplasma Infection**: Phytoplasma is a type of bacteria that can infect plants, including rhubarb. It can cause symptoms such as yellowing leaves, stunted growth, and splitting of petioles.\\n*   **Animal Damage**: Animals like rabbits, deer, or insects can damage rhubarb plants by eating the leaves or stems, which can lead to splitting of the petioles.\\n*   **Bacteria**: Bacterial infections can also cause damage to rhubarb plants, including splitting of the petioles.\\n\\n**Conclusion**\\n\\nBased on the analysis, it is clear that all the options listed \u001b[0m\u001b[32m(\u001b[0m\u001b[32mA\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Physiological problems, B\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Phytoplasma infection, D\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Animal damage, and E\u001b[0m\u001b[32m)\u001b[0m\u001b[32m Bacteria\u001b[0m\u001b[32m)\u001b[0m\u001b[32m could potentially cause the petioles of the rhubarb plant to split. Therefore, there is no single option that would not be a cause for the petioles splitting.\\n\\n**Answer**: C\u001b[0m\u001b[32m)\u001b[0m\u001b[32m I don't know and don't want to guess.\"\u001b[0m\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m\n",
               "\u001b[2;32m│   \u001b[0m\u001b[1m]\u001b[0m,\n",
               "\u001b[2;32m│   \u001b[0m\u001b[33mscores\u001b[0m=\u001b[1m{\u001b[0m\n",
@@ -815,7 +1005,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 8,
+      "execution_count": 5,
       "metadata": {
         "id": "HXmZf3Ymw-aX"
       },
@@ -823,39 +1013,33 @@
       "source": [
         "simpleqa_dataset_id = \"huggingface::simpleqa\"\n",
         "\n",
-        "_ = client.datasets.register(\n",
+        "register_dataset_response = client.datasets.register(\n",
+        "    purpose=\"eval/messages-answer\",\n",
+        "    source={\n",
+        "        \"type\": \"uri\",\n",
+        "        \"uri\": \"huggingface://datasets/llamastack/simpleqa?split=train\",\n",
+        "    },\n",
         "    dataset_id=simpleqa_dataset_id,\n",
-        "    provider_id=\"huggingface\",\n",
+        ")"
-        "    url={\"uri\": \"https://huggingface.co/datasets/llamastack/simpleqa\"},\n",
-        "    metadata={\n",
-        "        \"path\": \"llamastack/simpleqa\",\n",
-        "        \"split\": \"train\",\n",
-        "    },\n",
-        "    dataset_schema={\n",
-        "        \"input_query\": {\"type\": \"string\"},\n",
-        "        \"expected_answer\": {\"type\": \"string\"},\n",
-        "        \"chat_completion_input\": {\"type\": \"chat_completion_input\"},\n",
-        "    },\n",
-        ")\n"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 9,
+      "execution_count": 6,
       "metadata": {
         "id": "Gc8azb4Rxr5J"
       },
       "outputs": [],
       "source": [
-        "eval_rows = client.datasetio.get_rows_paginated(\n",
+        "eval_rows = client.datasets.iterrows(\n",
         "    dataset_id=simpleqa_dataset_id,\n",
-        "    rows_in_page=5,\n",
+        "    limit=5,\n",
-        ")\n"
+        ")"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 12,
+      "execution_count": 7,
       "metadata": {
         "colab": {
           "base_uri": "https://localhost:8080/",
@@ -876,7 +1060,7 @@
           "name": "stderr",
           "output_type": "stream",
           "text": [
-            "100%|██████████| 5/5 [00:31<00:00,  6.38s/it]\n"
+            "100%|██████████| 5/5 [00:13<00:00,  2.71s/it]\n"
           ]
         },
         {
@@ -889,14 +1073,14 @@
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">\"Radcliffe College was a women's liberal arts college in Cambridge, Massachusetts. However, it merged with Harvard University in 1977 and is now known as the Radcliffe Institute for Advanced Study at Harvard University.\"</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>,\n",
-              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'I do not have information on the Leipzig 1877 tournament.'</span><span style=\"font-weight: bold\">}</span>,\n",
+              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'I am unable to verify in whose honor the Leipzig 1877 tournament was organized.'</span><span style=\"font-weight: bold\">}</span>,\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">{</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'generated_answer'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">\"I am unable to verify what Empress Elizabeth of Austria's favorite sculpture depicted at her villa Achilleion at Corfu, according to Karl Küchler.\"</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"font-weight: bold\">}</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"font-weight: bold\">]</span>,\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   </span><span style=\"color: #808000; text-decoration-color: #808000\">scores</span>=<span style=\"font-weight: bold\">{</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   </span><span style=\"color: #008000; text-decoration-color: #008000\">'llm-as-judge::405b-simpleqa'</span>: <span style=\"color: #800080; text-decoration-color: #800080; font-weight: bold\">ScoringResult</span><span style=\"font-weight: bold\">(</span>\n",
-              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #808000; text-decoration-color: #808000\">aggregated_results</span>=<span style=\"font-weight: bold\">{}</span>,\n",
+              "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #808000; text-decoration-color: #808000\">aggregated_results</span>=<span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'categorical_count'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'categorical_count'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'A'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'C'</span>: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">4</span><span style=\"font-weight: bold\">}}}</span>,\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   </span><span style=\"color: #808000; text-decoration-color: #808000\">score_rows</span>=<span style=\"font-weight: bold\">[</span>\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'score'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'C'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'judge_feedback'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'C'</span><span style=\"font-weight: bold\">}</span>,\n",
               "<span style=\"color: #7fbf7f; text-decoration-color: #7fbf7f\">│   │   │   │   </span><span style=\"font-weight: bold\">{</span><span style=\"color: #008000; text-decoration-color: #008000\">'score'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'C'</span>, <span style=\"color: #008000; text-decoration-color: #008000\">'judge_feedback'</span>: <span style=\"color: #008000; text-decoration-color: #008000\">'C'</span><span style=\"font-weight: bold\">}</span>,\n",
@@ -917,14 +1101,14 @@
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\n",
               "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m\"Radcliffe College was a women's liberal arts college in Cambridge, Massachusetts. However, it merged with Harvard University in 1977 and is now known as the Radcliffe Institute for Advanced Study at Harvard University.\"\u001b[0m\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m,\n",
-              "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m'I do not have information on the Leipzig 1877 tournament.'\u001b[0m\u001b[1m}\u001b[0m,\n",
+              "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m'I am unable to verify in whose honor the Leipzig 1877 tournament was organized.'\u001b[0m\u001b[1m}\u001b[0m,\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m{\u001b[0m\n",
               "\u001b[2;32m│   │   │   \u001b[0m\u001b[32m'generated_answer'\u001b[0m: \u001b[32m\"I am unable to verify what Empress Elizabeth of Austria's favorite sculpture depicted at her villa Achilleion at Corfu, according to Karl Küchler.\"\u001b[0m\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[1m}\u001b[0m\n",
               "\u001b[2;32m│   \u001b[0m\u001b[1m]\u001b[0m,\n",
               "\u001b[2;32m│   \u001b[0m\u001b[33mscores\u001b[0m=\u001b[1m{\u001b[0m\n",
               "\u001b[2;32m│   │   \u001b[0m\u001b[32m'llm-as-judge::405b-simpleqa'\u001b[0m: \u001b[1;35mScoringResult\u001b[0m\u001b[1m(\u001b[0m\n",
-              "\u001b[2;32m│   │   │   \u001b[0m\u001b[33maggregated_results\u001b[0m=\u001b[1m{\u001b[0m\u001b[1m}\u001b[0m,\n",
+              "\u001b[2;32m│   │   │   \u001b[0m\u001b[33maggregated_results\u001b[0m=\u001b[1m{\u001b[0m\u001b[32m'categorical_count'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'categorical_count'\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m'A'\u001b[0m: \u001b[1;36m1\u001b[0m, \u001b[32m'C'\u001b[0m: \u001b[1;36m4\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m,\n",
               "\u001b[2;32m│   │   │   \u001b[0m\u001b[33mscore_rows\u001b[0m=\u001b[1m[\u001b[0m\n",
               "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[1m{\u001b[0m\u001b[32m'score'\u001b[0m: \u001b[32m'C'\u001b[0m, \u001b[32m'judge_feedback'\u001b[0m: \u001b[32m'C'\u001b[0m\u001b[1m}\u001b[0m,\n",
               "\u001b[2;32m│   │   │   │   \u001b[0m\u001b[1m{\u001b[0m\u001b[32m'score'\u001b[0m: \u001b[32m'C'\u001b[0m, \u001b[32m'judge_feedback'\u001b[0m: \u001b[32m'C'\u001b[0m\u001b[1m}\u001b[0m,\n",
@@ -957,7 +1141,7 @@
         "\n",
         "response = client.eval.evaluate_rows_alpha(\n",
         "    benchmark_id=\"meta-reference::simpleqa\",\n",
-        "    input_rows=eval_rows.rows,\n",
+        "    input_rows=eval_rows.data,\n",
         "    scoring_functions=[\"llm-as-judge::405b-simpleqa\"],\n",
         "    benchmark_config={\n",
         "        \"type\": \"benchmark\",\n",
@@ -1106,7 +1290,7 @@
         "\n",
         "response = client.eval.evaluate_rows_alpha(\n",
         "    benchmark_id=\"meta-reference::simpleqa\",\n",
-        "    input_rows=eval_rows.rows,\n",
+        "    input_rows=eval_rows.data,\n",
         "    scoring_functions=[\"llm-as-judge::405b-simpleqa\"],\n",
         "    benchmark_config={\n",
         "        \"type\": \"benchmark\",\n",

docs/openapi_generator/pyopenapi/generator.py

@@ -435,7 +435,7 @@ class Generator:
         )
         self.schema_builder = SchemaBuilder(schema_generator)
         self.responses = {}
-        
+
         # Create standard error responses
         self._create_standard_error_responses()
 
@@ -446,7 +446,7 @@ class Generator:
         """
         # Get the Error schema
         error_schema = self.schema_builder.classdef_to_ref(Error)
-        
+
         # Create standard error responses
         self.responses["BadRequest400"] = Response(
             description="The request was invalid or malformed",
@@ -457,11 +457,11 @@ class Generator:
                         "status": 400,
                         "title": "Bad Request",
                         "detail": "The request was invalid or malformed",
-                    }
+                    },
                 )
-            }
+            },
         )
-        
+
         self.responses["TooManyRequests429"] = Response(
             description="The client has sent too many requests in a given amount of time",
             content={
@@ -471,11 +471,11 @@ class Generator:
                         "status": 429,
                         "title": "Too Many Requests",
                         "detail": "You have exceeded the rate limit. Please try again later.",
-                    }
+                    },
                 )
-            }
+            },
         )
-        
+
         self.responses["InternalServerError500"] = Response(
             description="The server encountered an unexpected error",
             content={
@@ -485,11 +485,11 @@ class Generator:
                         "status": 500,
                         "title": "Internal Server Error",
                         "detail": "An unexpected error occurred. Our team has been notified.",
-                    }
+                    },
                 )
-            }
+            },
         )
-        
+
         # Add a default error response for any unhandled error cases
         self.responses["DefaultError"] = Response(
             description="An unexpected error occurred",
@@ -500,9 +500,9 @@ class Generator:
                         "status": 0,
                         "title": "Error",
                         "detail": "An unexpected error occurred",
-                    }
+                    },
                 )
-            }
+            },
         )
 
     def _build_type_tag(self, ref: str, schema: Schema) -> Tag:
@@ -547,11 +547,14 @@ class Generator:
             "SyntheticDataGeneration",
             "PostTraining",
             "BatchInference",
-            "Files",
         ]:
             op.defining_class.__name__ = f"{op.defining_class.__name__} (Coming Soon)"
             print(op.defining_class.__name__)
 
+        # TODO (xiyan): temporary fix for datasetio inner impl + datasets api
+        # if op.defining_class.__name__ in ["DatasetIO"]:
+        #     op.defining_class.__name__ = "Datasets"
+
         doc_string = parse_type(op.func_ref)
         doc_params = dict(
             (param.name, param.description) for param in doc_string.params.values()
@@ -598,7 +601,9 @@ class Generator:
 
         # data passed in request body as raw bytes cannot have request parameters
         if raw_bytes_request_body and op.request_params:
-            raise ValueError("Cannot have both raw bytes request body and request parameters")
+            raise ValueError(
+                "Cannot have both raw bytes request body and request parameters"
+            )
 
         # data passed in request body as raw bytes
         if raw_bytes_request_body:
@@ -719,7 +724,7 @@ class Generator:
             responses.update(response_builder.build_response(response_options))
 
         assert len(responses.keys()) > 0, f"No responses found for {op.name}"
-        
+
         # Add standard error response references
         if self.options.include_standard_error_responses:
             if "400" not in responses:
@@ -730,7 +735,7 @@ class Generator:
                 responses["500"] = ResponseRef("InternalServerError500")
             if "default" not in responses:
                 responses["default"] = ResponseRef("DefaultError")
-        
+
         if op.event_type is not None:
             builder = ContentBuilder(self.schema_builder)
             callbacks = {

docs/source/contributing/new_api_provider.md

@@ -6,7 +6,7 @@ This guide will walk you through the process of adding a new API provider to Lla
 - Begin by reviewing the [core concepts](../concepts/index.md) of Llama Stack and choose the API your provider belongs to (Inference, Safety, VectorIO, etc.)
 - Determine the provider type ({repopath}`Remote::llama_stack/providers/remote` or {repopath}`Inline::llama_stack/providers/inline`). Remote providers make requests to external services, while inline providers execute implementation locally.
 - Add your provider to the appropriate {repopath}`Registry::llama_stack/providers/registry/`. Specify pip dependencies necessary.
-- Update any distribution {repopath}`Templates::llama_stack/templates/` build.yaml and run.yaml files if they should include your provider by default. Run {repopath}`llama_stack/scripts/distro_codegen.py` if necessary. Note that `distro_codegen.py` will fail if the new provider causes any distribution template to attempt to import provider-specific dependencies. This usually means the distribution's `get_distribution_template()` code path should only import any necessary Config or model alias definitions from each provider and not the provider's actual implementation.
+- Update any distribution {repopath}`Templates::llama_stack/templates/` build.yaml and run.yaml files if they should include your provider by default. Run {repopath}`./scripts/distro_codegen.py` if necessary. Note that `distro_codegen.py` will fail if the new provider causes any distribution template to attempt to import provider-specific dependencies. This usually means the distribution's `get_distribution_template()` code path should only import any necessary Config or model alias definitions from each provider and not the provider's actual implementation.
 
 
 Here are some example PRs to help you get started:

docs/source/distributions/building_distro.md

@@ -185,8 +185,12 @@ llama stack build --config llama_stack/templates/ollama/build.yaml
 :::
 
 :::{tab-item} Building Container
-> [!TIP]
+
-> Podman is supported as an alternative to Docker. Set `CONTAINER_BINARY` to `podman` in your environment to use Podman.
+```{admonition} Podman Alternative
+:class: tip
+
+Podman is supported as an alternative to Docker. Set `CONTAINER_BINARY` to `podman` in your environment to use Podman.
+```
 
 To build a container image, you may start off from a template and use the `--image-type container` flag to specify `container` as the build image type.
 

docs/source/distributions/remote_hosted_distro/nvidia.md

@@ -6,13 +6,13 @@ The `llamastack/distribution-nvidia` distribution consists of the following prov
 | API | Provider(s) |
 |-----|-------------|
 | agents | `inline::meta-reference` |
-| datasetio | `remote::huggingface`, `inline::localfs` |
+| datasetio | `inline::localfs` |
 | eval | `inline::meta-reference` |
 | inference | `remote::nvidia` |
-| safety | `inline::llama-guard` |
+| safety | `remote::nvidia` |
-| scoring | `inline::basic`, `inline::llm-as-judge`, `inline::braintrust` |
+| scoring | `inline::basic` |
 | telemetry | `inline::meta-reference` |
-| tool_runtime | `remote::brave-search`, `remote::tavily-search`, `inline::code-interpreter`, `inline::rag-runtime`, `remote::model-context-protocol` |
+| tool_runtime | `inline::rag-runtime` |
 | vector_io | `inline::faiss` |
 
 
@@ -20,8 +20,10 @@ The `llamastack/distribution-nvidia` distribution consists of the following prov
 
 The following environment variables can be configured:
 
-- `LLAMASTACK_PORT`: Port for the Llama Stack distribution server (default: `5001`)
 - `NVIDIA_API_KEY`: NVIDIA API Key (default: ``)
+- `GUARDRAILS_SERVICE_URL`: URL for the NeMo Guardrails Service (default: `http://0.0.0.0:7331`)
+- `INFERENCE_MODEL`: Inference model (default: `Llama3.1-8B-Instruct`)
+- `SAFETY_MODEL`: Name of the model to use for safety (default: `meta/llama-3.1-8b-instruct`)
 
 ### Models
 

docs/source/playground/index.md

@@ -92,6 +92,8 @@ Interactive pages for users to play with and explore Llama Stack API capabilitie
 
 ## Starting the Llama Stack Playground
 
+### Llama CLI
+
 To start the Llama Stack Playground, run the following commands:
 
 1. Start up the Llama Stack API server
@@ -107,3 +109,28 @@ cd llama_stack/distribution/ui
 pip install -r requirements.txt
 streamlit run app.py
 ```
+
+### Docker
+
+Playground can also be started in a docker image:
+
+```sh
+export LLAMA_STACK_URL=http://localhost:11434
+
+docker run \
+  -p 8501:8501 \
+  -e LLAMA_STACK_ENDPOINT=$LLAMA_STACK_URL \
+  quay.io/jland/llama-stack-playground
+```
+
+## Configurable Environment Variables
+
+## Environment Variables
+
+| Environment Variable       | Description                        | Default Value             |
+|----------------------------|------------------------------------|---------------------------|
+| LLAMA_STACK_ENDPOINT       | The endpoint for the Llama Stack   | http://localhost:8321     |
+| FIREWORKS_API_KEY          | API key for Fireworks provider     | (empty string)            |
+| TOGETHER_API_KEY           | API key for Together provider      | (empty string)            |
+| SAMBANOVA_API_KEY          | API key for SambaNova provider     | (empty string)            |
+| OPENAI_API_KEY             | API key for OpenAI provider        | (empty string)            |

docs/source/references/evals_reference/index.md

@@ -114,23 +114,17 @@ pprint(response)
 simpleqa_dataset_id = "huggingface::simpleqa"
 
 _ = client.datasets.register(
+    purpose="eval/messages-answer",
+    source={
+        "type": "uri",
+        "uri": "huggingface://datasets/llamastack/simpleqa?split=train",
+    },
     dataset_id=simpleqa_dataset_id,
-    provider_id="huggingface",
-    url={"uri": "https://huggingface.co/datasets/llamastack/simpleqa"},
-    metadata={
-        "path": "llamastack/simpleqa",
-        "split": "train",
-    },
-    dataset_schema={
-        "input_query": {"type": "string"},
-        "expected_answer": {"type": "string"},
-        "chat_completion_input": {"type": "chat_completion_input"},
-    },
 )
 
-eval_rows = client.datasetio.get_rows_paginated(
+eval_rows = client.datasets.iterrows(
     dataset_id=simpleqa_dataset_id,
-    rows_in_page=5,
+    limit=5,
 )
 ```
 
@@ -143,7 +137,7 @@ client.benchmarks.register(
 
 response = client.eval.evaluate_rows(
     benchmark_id="meta-reference::simpleqa",
-    input_rows=eval_rows.rows,
+    input_rows=eval_rows.data,
     scoring_functions=["llm-as-judge::405b-simpleqa"],
     benchmark_config={
         "eval_candidate": {
@@ -191,7 +185,7 @@ agent_config = {
 
 response = client.eval.evaluate_rows(
     benchmark_id="meta-reference::simpleqa",
-    input_rows=eval_rows.rows,
+    input_rows=eval_rows.data,
     scoring_functions=["llm-as-judge::405b-simpleqa"],
     benchmark_config={
         "eval_candidate": {

docs/source/references/llama_stack_client_cli_reference.md

@@ -6,17 +6,32 @@ The `llama-stack-client` CLI allows you to query information about the distribut
 
 ### `llama-stack-client`
 ```bash
-llama-stack-client -h
+llama-stack-client
+Usage: llama-stack-client [OPTIONS] COMMAND [ARGS]...
 
-usage: llama-stack-client [-h] {models,memory_banks,shields} ...
+  Welcome to the LlamaStackClient CLI
 
-Welcome to the LlamaStackClient CLI
+Options:
+  --version        Show the version and exit.
+  --endpoint TEXT  Llama Stack distribution endpoint
+  --api-key TEXT   Llama Stack distribution API key
+  --config TEXT    Path to config file
+  --help           Show this message and exit.
 
-options:
+Commands:
-  -h, --help            show this help message and exit
+  configure          Configure Llama Stack Client CLI.
-
+  datasets           Manage datasets.
-subcommands:
+  eval               Run evaluation tasks.
-  {models,memory_banks,shields}
+  eval_tasks         Manage evaluation tasks.
+  inference          Inference (chat).
+  inspect            Inspect server configuration.
+  models             Manage GenAI models.
+  post_training      Post-training.
+  providers          Manage API providers.
+  scoring_functions  Manage scoring functions.
+  shields            Manage safety shield services.
+  toolgroups         Manage available tool groups.
+  vector_dbs         Manage vector databases.
 ```
 
 ### `llama-stack-client configure`
@@ -127,11 +142,11 @@ llama-stack-client vector_dbs list
 llama-stack-client vector_dbs register <vector-db-id> [--provider-id <provider-id>] [--provider-vector-db-id <provider-vector-db-id>] [--embedding-model <embedding-model>] [--embedding-dimension <embedding-dimension>]
 ```
 
-Options:
+Optional arguments:
-- `--provider-id`: Optional. Provider ID for the vector db
+- `--provider-id`: Provider ID for the vector db
-- `--provider-vector-db-id`: Optional. Provider's vector db ID
+- `--provider-vector-db-id`: Provider's vector db ID
-- `--embedding-model`: Optional. Embedding model to use. Default: "all-MiniLM-L6-v2"
+- `--embedding-model`: Embedding model to use. Default: "all-MiniLM-L6-v2"
-- `--embedding-dimension`: Optional. Dimension of embeddings. Default: 384
+- `--embedding-dimension`: Dimension of embeddings. Default: 384
 
 ### `llama-stack-client vector_dbs unregister`
 ```bash
@@ -157,11 +172,13 @@ llama-stack-client shields list
 llama-stack-client shields register --shield-id <shield-id> [--provider-id <provider-id>] [--provider-shield-id <provider-shield-id>] [--params <params>]
 ```
 
-Options:
+Required arguments:
-- `--shield-id`: Required. ID of the shield
+- `--shield-id`: ID of the shield
-- `--provider-id`: Optional. Provider ID for the shield
+
-- `--provider-shield-id`: Optional. Provider's shield ID
+Optional arguments:
-- `--params`: Optional. JSON configuration parameters for the shield
+- `--provider-id`: Provider ID for the shield
+- `--provider-shield-id`: Provider's shield ID
+- `--params`: JSON configuration parameters for the shield
 
 ## Eval Task Management
 
@@ -175,13 +192,15 @@ llama-stack-client benchmarks list
 llama-stack-client benchmarks register --eval-task-id <eval-task-id> --dataset-id <dataset-id> --scoring-functions <function1> [<function2> ...] [--provider-id <provider-id>] [--provider-eval-task-id <provider-eval-task-id>] [--metadata <metadata>]
 ```
 
-Options:
+Required arguments:
-- `--eval-task-id`: Required. ID of the eval task
+- `--eval-task-id`: ID of the eval task
-- `--dataset-id`: Required. ID of the dataset to evaluate
+- `--dataset-id`: ID of the dataset to evaluate
-- `--scoring-functions`: Required. One or more scoring functions to use for evaluation
+- `--scoring-functions`: One or more scoring functions to use for evaluation
-- `--provider-id`: Optional. Provider ID for the eval task
+
-- `--provider-eval-task-id`: Optional. Provider's eval task ID
+Optional arguments:
-- `--metadata`: Optional. Metadata for the eval task in JSON format
+- `--provider-id`: Provider ID for the eval task
+- `--provider-eval-task-id`: Provider's eval task ID
+- `--metadata`: Metadata for the eval task in JSON format
 
 ## Eval execution
 ### `llama-stack-client eval run-benchmark`
@@ -189,11 +208,13 @@ Options:
 llama-stack-client eval run-benchmark <eval-task-id1> [<eval-task-id2> ...] --eval-task-config <config-file> --output-dir <output-dir> [--num-examples <num>] [--visualize]
 ```
 
-Options:
+Required arguments:
-- `--eval-task-config`: Required. Path to the eval task config file in JSON format
+- `--eval-task-config`: Path to the eval task config file in JSON format
-- `--output-dir`: Required. Path to the directory where evaluation results will be saved
+- `--output-dir`: Path to the directory where evaluation results will be saved
-- `--num-examples`: Optional. Number of examples to evaluate (useful for debugging)
+
-- `--visualize`: Optional flag. If set, visualizes evaluation results after completion
+Optional arguments:
+- `--num-examples`: Number of examples to evaluate (useful for debugging)
+- `--visualize`: If set, visualizes evaluation results after completion
 
 Example benchmark_config.json:
 ```json
@@ -214,11 +235,13 @@ Example benchmark_config.json:
 llama-stack-client eval run-scoring <eval-task-id> --eval-task-config <config-file> --output-dir <output-dir> [--num-examples <num>] [--visualize]
 ```
 
-Options:
+Required arguments:
-- `--eval-task-config`: Required. Path to the eval task config file in JSON format
+- `--eval-task-config`: Path to the eval task config file in JSON format
-- `--output-dir`: Required. Path to the directory where scoring results will be saved
+- `--output-dir`: Path to the directory where scoring results will be saved
-- `--num-examples`: Optional. Number of examples to evaluate (useful for debugging)
+
-- `--visualize`: Optional flag. If set, visualizes scoring results after completion
+Optional arguments:
+- `--num-examples`: Number of examples to evaluate (useful for debugging)
+- `--visualize`: If set, visualizes scoring results after completion
 
 ## Tool Group Management
 
@@ -230,11 +253,11 @@ llama-stack-client toolgroups list
 +---------------------------+------------------+------+---------------+
 | identifier                | provider_id      | args | mcp_endpoint  |
 +===========================+==================+======+===============+
-| builtin::code_interpreter | code-interpreter | None | None         |
+| builtin::code_interpreter | code-interpreter | None | None          |
 +---------------------------+------------------+------+---------------+
-| builtin::rag             | rag-runtime      | None | None         |
+| builtin::rag              | rag-runtime      | None | None          |
 +---------------------------+------------------+------+---------------+
-| builtin::websearch       | tavily-search    | None | None         |
+| builtin::websearch        | tavily-search    | None | None          |
 +---------------------------+------------------+------+---------------+
 ```
 
@@ -250,11 +273,11 @@ Shows detailed information about a specific toolgroup. If the toolgroup is not f
 llama-stack-client toolgroups register <toolgroup_id> [--provider-id <provider-id>] [--provider-toolgroup-id <provider-toolgroup-id>] [--mcp-config <mcp-config>] [--args <args>]
 ```
 
-Options:
+Optional arguments:
-- `--provider-id`: Optional. Provider ID for the toolgroup
+- `--provider-id`: Provider ID for the toolgroup
-- `--provider-toolgroup-id`: Optional. Provider's toolgroup ID
+- `--provider-toolgroup-id`: Provider's toolgroup ID
-- `--mcp-config`: Optional. JSON configuration for the MCP endpoint
+- `--mcp-config`: JSON configuration for the MCP endpoint
-- `--args`: Optional. JSON arguments for the toolgroup
+- `--args`: JSON arguments for the toolgroup
 
 ### `llama-stack-client toolgroups unregister`
 ```bash

llama_stack/apis/datasetio/datasetio.py

@@ -13,19 +13,16 @@ from llama_stack.schema_utils import json_schema_type, webmethod
 
 
 @json_schema_type
-class PaginatedRowsResult(BaseModel):
+class IterrowsResponse(BaseModel):
     """
     A paginated list of rows from a dataset.
 
-    :param rows: The rows in the current page.
+    :param data: The rows in the current page.
-    :param total_count: The total number of rows in the dataset.
+    :param next_start_index: Index into dataset for the first row in the next page. None if there are no more rows.
-    :param next_page_token: The token to get the next page of rows.
     """
 
-    # the rows obey the DatasetSchema for the given dataset
+    data: List[Dict[str, Any]]
-    rows: List[Dict[str, Any]]
+    next_start_index: Optional[int] = None
-    total_count: int
-    next_page_token: Optional[str] = None
 
 
 class DatasetStore(Protocol):
@@ -37,22 +34,21 @@ class DatasetIO(Protocol):
     # keeping for aligning with inference/safety, but this is not used
     dataset_store: DatasetStore
 
-    @webmethod(route="/datasetio/rows", method="GET")
+    # TODO(xiyan): there's a flakiness here where setting route to "/datasets/" here will not result in proper routing
-    async def get_rows_paginated(
+    @webmethod(route="/datasetio/iterrows/{dataset_id:path}", method="GET")
+    async def iterrows(
         self,
         dataset_id: str,
-        rows_in_page: int,
+        start_index: Optional[int] = None,
-        page_token: Optional[str] = None,
+        limit: Optional[int] = None,
-        filter_condition: Optional[str] = None,
+    ) -> IterrowsResponse:
-    ) -> PaginatedRowsResult:
+        """Get a paginated list of rows from a dataset. Uses cursor-based pagination.
-        """Get a paginated list of rows from a dataset.
 
         :param dataset_id: The ID of the dataset to get the rows from.
-        :param rows_in_page: The number of rows to get per page.
+        :param start_index: Index into dataset for the first row to get. Get all rows if None.
-        :param page_token: The token to get the next page of rows.
+        :param limit: The number of rows to get.
-        :param filter_condition: (Optional) A condition to filter the rows by.
         """
         ...
 
-    @webmethod(route="/datasetio/rows", method="POST")
+    @webmethod(route="/datasetio/append-rows/{dataset_id:path}", method="POST")
     async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None: ...

llama_stack/apis/datasets/datasets.py

@@ -4,19 +4,102 @@
 # 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, Dict, List, Literal, Optional, Protocol
+from enum import Enum
+from typing import Annotated, Any, Dict, List, Literal, Optional, Protocol, Union
 
 from pydantic import BaseModel, Field
 
-from llama_stack.apis.common.content_types import URL
-from llama_stack.apis.common.type_system import ParamType
 from llama_stack.apis.resource import Resource, ResourceType
-from llama_stack.schema_utils import json_schema_type, webmethod
+from llama_stack.schema_utils import json_schema_type, register_schema, webmethod
+
+
+class DatasetPurpose(str, Enum):
+    """
+    Purpose of the dataset. Each purpose has a required input data schema.
+
+    :cvar post-training/messages: The dataset contains messages used for post-training.
+        {
+            "messages": [
+                {"role": "user", "content": "Hello, world!"},
+                {"role": "assistant", "content": "Hello, world!"},
+            ]
+        }
+    :cvar eval/question-answer: The dataset contains a question column and an answer column.
+        {
+            "question": "What is the capital of France?",
+            "answer": "Paris"
+        }
+    :cvar eval/messages-answer: The dataset contains a messages column with list of messages and an answer column.
+        {
+            "messages": [
+                {"role": "user", "content": "Hello, my name is John Doe."},
+                {"role": "assistant", "content": "Hello, John Doe. How can I help you today?"},
+                {"role": "user", "content": "What's my name?"},
+            ],
+            "answer": "John Doe"
+        }
+    """
+
+    post_training_messages = "post-training/messages"
+    eval_question_answer = "eval/question-answer"
+    eval_messages_answer = "eval/messages-answer"
+
+    # TODO: add more schemas here
+
+
+class DatasetType(Enum):
+    """
+    Type of the dataset source.
+    :cvar uri: The dataset can be obtained from a URI.
+    :cvar rows: The dataset is stored in rows.
+    """
+
+    uri = "uri"
+    rows = "rows"
+
+
+@json_schema_type
+class URIDataSource(BaseModel):
+    """A dataset that can be obtained from a URI.
+    :param uri: The dataset can be obtained from a URI. E.g.
+        - "https://mywebsite.com/mydata.jsonl"
+        - "lsfs://mydata.jsonl"
+        - "data:csv;base64,{base64_content}"
+    """
+
+    type: Literal["uri"] = "uri"
+    uri: str
+
+
+@json_schema_type
+class RowsDataSource(BaseModel):
+    """A dataset stored in rows.
+    :param rows: The dataset is stored in rows. E.g.
+        - [
+            {"messages": [{"role": "user", "content": "Hello, world!"}, {"role": "assistant", "content": "Hello, world!"}]}
+        ]
+    """
+
+    type: Literal["rows"] = "rows"
+    rows: List[Dict[str, Any]]
+
+
+DataSource = register_schema(
+    Annotated[
+        Union[URIDataSource, RowsDataSource],
+        Field(discriminator="type"),
+    ],
+    name="DataSource",
+)
 
 
 class CommonDatasetFields(BaseModel):
-    dataset_schema: Dict[str, ParamType]
+    """
-    url: URL
+    Common fields for a dataset.
+    """
+
+    purpose: DatasetPurpose
+    source: DataSource
     metadata: Dict[str, Any] = Field(
         default_factory=dict,
         description="Any additional metadata for this dataset",
@@ -50,13 +133,69 @@ class Datasets(Protocol):
     @webmethod(route="/datasets", method="POST")
     async def register_dataset(
         self,
-        dataset_id: str,
+        purpose: DatasetPurpose,
-        dataset_schema: Dict[str, ParamType],
+        source: DataSource,
-        url: URL,
-        provider_dataset_id: Optional[str] = None,
-        provider_id: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
-    ) -> None: ...
+        dataset_id: Optional[str] = None,
+    ) -> Dataset:
+        """
+        Register a new dataset.
+
+        :param purpose: The purpose of the dataset. One of
+            - "post-training/messages": The dataset contains a messages column with list of messages for post-training.
+                {
+                    "messages": [
+                        {"role": "user", "content": "Hello, world!"},
+                        {"role": "assistant", "content": "Hello, world!"},
+                    ]
+                }
+            - "eval/question-answer": The dataset contains a question column and an answer column for evaluation.
+                {
+                    "question": "What is the capital of France?",
+                    "answer": "Paris"
+                }
+            - "eval/messages-answer": The dataset contains a messages column with list of messages and an answer column for evaluation.
+                {
+                    "messages": [
+                        {"role": "user", "content": "Hello, my name is John Doe."},
+                        {"role": "assistant", "content": "Hello, John Doe. How can I help you today?"},
+                        {"role": "user", "content": "What's my name?"},
+                    ],
+                    "answer": "John Doe"
+                }
+        :param source: The data source of the dataset. Ensure that the data source schema is compatible with the purpose of the dataset. Examples:
+           - {
+               "type": "uri",
+               "uri": "https://mywebsite.com/mydata.jsonl"
+           }
+           - {
+               "type": "uri",
+               "uri": "lsfs://mydata.jsonl"
+           }
+           - {
+               "type": "uri",
+               "uri": "data:csv;base64,{base64_content}"
+           }
+           - {
+               "type": "uri",
+               "uri": "huggingface://llamastack/simpleqa?split=train"
+           }
+           - {
+               "type": "rows",
+               "rows": [
+                   {
+                       "messages": [
+                           {"role": "user", "content": "Hello, world!"},
+                           {"role": "assistant", "content": "Hello, world!"},
+                       ]
+                   }
+               ]
+           }
+        :param metadata: The metadata for the dataset.
+           - E.g. {"description": "My dataset"}
+        :param dataset_id: The ID of the dataset. If not provided, an ID will be generated.
+        """
+        ...
 
     @webmethod(route="/datasets/{dataset_id:path}", method="GET")
     async def get_dataset(

llama_stack/cli/stack/_build.py

@@ -38,7 +38,7 @@ from llama_stack.distribution.distribution import get_provider_registry
 from llama_stack.distribution.resolver import InvalidProviderError
 from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR
 from llama_stack.distribution.utils.dynamic import instantiate_class_type
-from llama_stack.distribution.utils.exec import formulate_run_args, run_with_pty
+from llama_stack.distribution.utils.exec import formulate_run_args, run_command
 from llama_stack.distribution.utils.image_types import LlamaStackImageType
 from llama_stack.providers.datatypes import Api
 
@@ -213,7 +213,7 @@ def run_stack_build_command(args: argparse.Namespace) -> None:
         config = parse_and_maybe_upgrade_config(config_dict)
         run_args = formulate_run_args(args.image_type, args.image_name, config, args.template)
         run_args.extend([run_config, str(os.getenv("LLAMA_STACK_PORT", 8321))])
-        run_with_pty(run_args)
+        run_command(run_args)
 
 
 def _generate_run_config(

llama_stack/cli/stack/run.py

@@ -82,7 +82,7 @@ class StackRun(Subcommand):
 
         from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
         from llama_stack.distribution.utils.config_dirs import DISTRIBS_BASE_DIR
-        from llama_stack.distribution.utils.exec import formulate_run_args, run_with_pty
+        from llama_stack.distribution.utils.exec import formulate_run_args, run_command
 
         config_file = Path(args.config)
         has_yaml_suffix = args.config.endswith(".yaml")
@@ -136,4 +136,4 @@ class StackRun(Subcommand):
 
         if args.tls_keyfile and args.tls_certfile:
             run_args.extend(["--tls-keyfile", args.tls_keyfile, "--tls-certfile", args.tls_certfile])
-        run_with_pty(run_args)
+        run_command(run_args)

llama_stack/distribution/build.py

@@ -6,7 +6,6 @@
 
 import importlib.resources
 import logging
-import sys
 from pathlib import Path
 from typing import Dict, List
 
@@ -15,7 +14,7 @@ from termcolor import cprint
 
 from llama_stack.distribution.datatypes import BuildConfig, Provider
 from llama_stack.distribution.distribution import get_provider_registry
-from llama_stack.distribution.utils.exec import run_command, run_with_pty
+from llama_stack.distribution.utils.exec import run_command
 from llama_stack.distribution.utils.image_types import LlamaStackImageType
 from llama_stack.providers.datatypes import Api
 
@@ -123,11 +122,7 @@ def build_image(
     if special_deps:
         args.append("#".join(special_deps))
 
-    is_terminal = sys.stdin.isatty()
+    return_code = run_command(args)
-    if is_terminal:
-        return_code = run_with_pty(args)
-    else:
-        return_code = run_command(args)
 
     if return_code != 0:
         log.error(

llama_stack/distribution/build_container.sh

@@ -43,7 +43,7 @@ RED='\033[0;31m'
 NC='\033[0m' # No Color
 
 CONTAINER_BINARY=${CONTAINER_BINARY:-docker}
-CONTAINER_OPTS=${CONTAINER_OPTS:-}
+CONTAINER_OPTS=${CONTAINER_OPTS:---progress=plain}
 
 TEMP_DIR=$(mktemp -d)
 
@@ -253,8 +253,7 @@ $CONTAINER_BINARY build \
   "${CLI_ARGS[@]}" \
   -t "$image_tag" \
   -f "$TEMP_DIR/Containerfile" \
-  "." \
+  "."
-  --progress=plain
 
 # clean up tmp/configs
 set +x

llama_stack/distribution/providers.py

@@ -8,10 +8,13 @@
 from pydantic import BaseModel
 
 from llama_stack.apis.providers import ListProvidersResponse, ProviderInfo, Providers
+from llama_stack.log import get_logger
 
 from .datatypes import StackRunConfig
 from .stack import redact_sensitive_fields
 
+logger = get_logger(name=__name__, category="core")
+
 
 class ProviderImplConfig(BaseModel):
     run_config: StackRunConfig
@@ -31,6 +34,10 @@ class ProviderImpl(Providers):
     async def initialize(self) -> None:
         pass
 
+    async def shutdown(self) -> None:
+        logger.debug("ProviderImpl.shutdown")
+        pass
+
     async def list_providers(self) -> ListProvidersResponse:
         run_config = self.config.run_config
         safe_config = StackRunConfig(**redact_sensitive_fields(run_config.model_dump()))

llama_stack/distribution/routers/routers.py

@@ -12,7 +12,8 @@ from llama_stack.apis.common.content_types import (
     InterleavedContent,
     InterleavedContentItem,
 )
-from llama_stack.apis.datasetio import DatasetIO, PaginatedRowsResult
+from llama_stack.apis.datasetio import DatasetIO, IterrowsResponse
+from llama_stack.apis.datasets import DatasetPurpose, DataSource
 from llama_stack.apis.eval import (
     BenchmarkConfig,
     Eval,
@@ -160,7 +161,11 @@ class InferenceRouter(Inference):
         await self.routing_table.register_model(model_id, provider_model_id, provider_id, metadata, model_type)
 
     def _construct_metrics(
-        self, prompt_tokens: int, completion_tokens: int, total_tokens: int, model: Model
+        self,
+        prompt_tokens: int,
+        completion_tokens: int,
+        total_tokens: int,
+        model: Model,
     ) -> List[MetricEvent]:
         """Constructs a list of MetricEvent objects containing token usage metrics.
 
@@ -298,7 +303,12 @@ class InferenceRouter(Inference):
                             completion_text += chunk.event.delta.text
                     if chunk.event.event_type == ChatCompletionResponseEventType.complete:
                         completion_tokens = await self._count_tokens(
-                            [CompletionMessage(content=completion_text, stop_reason=StopReason.end_of_turn)],
+                            [
+                                CompletionMessage(
+                                    content=completion_text,
+                                    stop_reason=StopReason.end_of_turn,
+                                )
+                            ],
                             tool_config.tool_prompt_format,
                         )
                         total_tokens = (prompt_tokens or 0) + (completion_tokens or 0)
@@ -471,21 +481,36 @@ class DatasetIORouter(DatasetIO):
         logger.debug("DatasetIORouter.shutdown")
         pass
 
-    async def get_rows_paginated(
+    async def register_dataset(
+        self,
+        purpose: DatasetPurpose,
+        source: DataSource,
+        metadata: Optional[Dict[str, Any]] = None,
+        dataset_id: Optional[str] = None,
+    ) -> None:
+        logger.debug(
+            f"DatasetIORouter.register_dataset: {purpose=} {source=} {metadata=} {dataset_id=}",
+        )
+        await self.routing_table.register_dataset(
+            purpose=purpose,
+            source=source,
+            metadata=metadata,
+            dataset_id=dataset_id,
+        )
+
+    async def iterrows(
         self,
         dataset_id: str,
-        rows_in_page: int,
+        start_index: Optional[int] = None,
-        page_token: Optional[str] = None,
+        limit: Optional[int] = None,
-        filter_condition: Optional[str] = None,
+    ) -> IterrowsResponse:
-    ) -> PaginatedRowsResult:
         logger.debug(
-            f"DatasetIORouter.get_rows_paginated: {dataset_id}, rows_in_page={rows_in_page}",
+            f"DatasetIORouter.iterrows: {dataset_id}, {start_index=} {limit=}",
         )
-        return await self.routing_table.get_provider_impl(dataset_id).get_rows_paginated(
+        return await self.routing_table.get_provider_impl(dataset_id).iterrows(
             dataset_id=dataset_id,
-            rows_in_page=rows_in_page,
+            start_index=start_index,
-            page_token=page_token,
+            limit=limit,
-            filter_condition=filter_condition,
         )
 
     async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:

llama_stack/distribution/routers/routing_tables.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 import logging
+import uuid
 from typing import Any, Dict, List, Optional
 
 from pydantic import TypeAdapter
@@ -12,7 +13,14 @@ from pydantic import TypeAdapter
 from llama_stack.apis.benchmarks import Benchmark, Benchmarks, ListBenchmarksResponse
 from llama_stack.apis.common.content_types import URL
 from llama_stack.apis.common.type_system import ParamType
-from llama_stack.apis.datasets import Dataset, Datasets, ListDatasetsResponse
+from llama_stack.apis.datasets import (
+    Dataset,
+    DatasetPurpose,
+    Datasets,
+    DatasetType,
+    DataSource,
+    ListDatasetsResponse,
+)
 from llama_stack.apis.models import ListModelsResponse, Model, Models, ModelType
 from llama_stack.apis.resource import ResourceType
 from llama_stack.apis.scoring_functions import (
@@ -352,34 +360,42 @@ class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):
 
     async def register_dataset(
         self,
-        dataset_id: str,
+        purpose: DatasetPurpose,
-        dataset_schema: Dict[str, ParamType],
+        source: DataSource,
-        url: URL,
-        provider_dataset_id: Optional[str] = None,
-        provider_id: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
-    ) -> None:
+        dataset_id: Optional[str] = None,
-        if provider_dataset_id is None:
+    ) -> Dataset:
-            provider_dataset_id = dataset_id
+        if not dataset_id:
-        if provider_id is None:
+            dataset_id = f"dataset-{str(uuid.uuid4())}"
-            # If provider_id not specified, use the only provider if it supports this dataset
+
-            if len(self.impls_by_provider_id) == 1:
+        provider_dataset_id = dataset_id
-                provider_id = list(self.impls_by_provider_id.keys())[0]
+
+        # infer provider from source
+        if source.type == DatasetType.rows.value:
+            provider_id = "localfs"
+        elif source.type == DatasetType.uri.value:
+            # infer provider from uri
+            if source.uri.startswith("huggingface"):
+                provider_id = "huggingface"
             else:
-                raise ValueError(
+                provider_id = "localfs"
-                    f"No provider specified and multiple providers available. Please specify a provider_id. Available providers: {self.impls_by_provider_id.keys()}"
+        else:
-                )
+            raise ValueError(f"Unknown data source type: {source.type}")
+
         if metadata is None:
             metadata = {}
+
         dataset = Dataset(
             identifier=dataset_id,
             provider_resource_id=provider_dataset_id,
             provider_id=provider_id,
-            dataset_schema=dataset_schema,
+            purpose=purpose,
-            url=url,
+            source=source,
             metadata=metadata,
         )
+
         await self.register_object(dataset)
+        return dataset
 
     async def unregister_dataset(self, dataset_id: str) -> None:
         dataset = await self.get_dataset(dataset_id)

llama_stack/distribution/ui/Containerfile

@@ -0,0 +1,11 @@
+# More info on playground configuration can be found here:
+# https://llama-stack.readthedocs.io/en/latest/playground
+
+FROM python:3.9-slim
+WORKDIR /app
+COPY . /app/
+RUN /usr/local/bin/python -m pip install --upgrade pip && \
+    /usr/local/bin/pip3 install -r requirements.txt
+EXPOSE 8501
+
+ENTRYPOINT ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]

llama_stack/distribution/ui/README.md

@@ -40,3 +40,13 @@ cd llama_stack/distribution/ui
 pip install -r requirements.txt
 streamlit run app.py
 ```
+
+## Environment Variables
+
+| Environment Variable       | Description                        | Default Value             |
+|----------------------------|------------------------------------|---------------------------|
+| LLAMA_STACK_ENDPOINT       | The endpoint for the Llama Stack   | http://localhost:8321     |
+| FIREWORKS_API_KEY          | API key for Fireworks provider     | (empty string)            |
+| TOGETHER_API_KEY           | API key for Together provider      | (empty string)            |
+| SAMBANOVA_API_KEY          | API key for SambaNova provider     | (empty string)            |
+| OPENAI_API_KEY             | API key for OpenAI provider        | (empty string)            |

llama_stack/distribution/ui/page/evaluations/native_eval.py

@@ -166,11 +166,10 @@ def run_evaluation_3():
     eval_candidate = st.session_state["eval_candidate"]
 
     dataset_id = benchmarks[selected_benchmark].dataset_id
-    rows = llama_stack_api.client.datasetio.get_rows_paginated(
+    rows = llama_stack_api.client.datasets.iterrows(
         dataset_id=dataset_id,
-        rows_in_page=-1,
     )
-    total_rows = len(rows.rows)
+    total_rows = len(rows.data)
     # Add number of examples control
     num_rows = st.number_input(
         "Number of Examples to Evaluate",
@@ -195,7 +194,7 @@ def run_evaluation_3():
     if st.button("Run Evaluation"):
         progress_text = "Running evaluation..."
         progress_bar = st.progress(0, text=progress_text)
-        rows = rows.rows
+        rows = rows.data
         if num_rows < total_rows:
             rows = rows[:num_rows]
 

llama_stack/distribution/utils/exec.py

@@ -4,13 +4,10 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import errno
 import logging
 import os
-import select
 import signal
 import subprocess
-import sys
 
 from termcolor import cprint
 
@@ -88,13 +85,6 @@ def formulate_run_args(image_type, image_name, config, template_name) -> list:
     return run_args
 
 
-def run_with_pty(command):
-    if sys.platform.startswith("win"):
-        return _run_with_pty_win(command)
-    else:
-        return _run_with_pty_unix(command)
-
-
 def in_notebook():
     try:
         from IPython import get_ipython
@@ -108,19 +98,19 @@ def in_notebook():
     return True
 
 
-# run a command in a pseudo-terminal, with interrupt handling,
+def run_command(command: list[str]) -> int:
-# useful when you want to run interactive things
+    """
-def _run_with_pty_unix(command):
+    Run a command with interrupt handling and output capture.
-    import pty
+    Uses subprocess.run with direct stream piping for better performance.
-    import termios
 
-    master, slave = pty.openpty()
+    Args:
+        command (list): The command to run.
 
-    old_settings = termios.tcgetattr(sys.stdin)
+    Returns:
+        int: The return code of the command.
+    """
     original_sigint = signal.getsignal(signal.SIGINT)
-
     ctrl_c_pressed = False
-    process = None
 
     def sigint_handler(signum, frame):
         nonlocal ctrl_c_pressed
@@ -131,106 +121,19 @@ def _run_with_pty_unix(command):
         # Set up the signal handler
         signal.signal(signal.SIGINT, sigint_handler)
 
-        new_settings = termios.tcgetattr(sys.stdin)
+        # Run the command with stdout/stderr piped directly to system streams
-        new_settings[3] = new_settings[3] & ~termios.ECHO  # Disable echo
+        result = subprocess.run(
-        new_settings[3] = new_settings[3] & ~termios.ICANON  # Disable canonical mode
-        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, new_settings)
-
-        process = subprocess.Popen(
             command,
-            stdin=slave,
+            text=True,
-            stdout=slave,
+            check=False,
-            stderr=slave,
-            universal_newlines=True,
-            preexec_fn=os.setsid,
         )
-
+        return result.returncode
-        # Close the slave file descriptor as it's now owned by the subprocess
+    except subprocess.SubprocessError as e:
-        os.close(slave)
+        log.error(f"Subprocess error: {e}")
-
+        return 1
-        def handle_io():
-            while not ctrl_c_pressed:
-                try:
-                    rlist, _, _ = select.select([sys.stdin, master], [], [], 0.1)
-
-                    if sys.stdin in rlist:
-                        data = os.read(sys.stdin.fileno(), 1024)
-                        if not data:
-                            break
-                        os.write(master, data)
-
-                    if master in rlist:
-                        data = os.read(master, 1024)
-                        if not data:
-                            break
-                        sys.stdout.buffer.write(data)
-                        sys.stdout.flush()
-
-                except KeyboardInterrupt:
-                    # This will be raised when Ctrl+C is pressed
-                    break
-
-                if process.poll() is not None:
-                    break
-
-        handle_io()
-    except (EOFError, KeyboardInterrupt):
-        pass
-    except OSError as e:
-        if e.errno != errno.EIO:
-            raise
-    finally:
-        # Clean up
-        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)
-        signal.signal(signal.SIGINT, original_sigint)
-
-        os.close(master)
-        if process and process.poll() is None:
-            process.terminate()
-            process.wait()
-
-    return process.returncode
-
-
-# run a command in a pseudo-terminal in windows, with interrupt handling,
-def _run_with_pty_win(command):
-    """
-    Runs a command with interactive support using subprocess directly.
-    """
-    try:
-        # For shell scripts on Windows, use appropriate shell
-        if isinstance(command, (list, tuple)):
-            if command[0].endswith(".sh"):
-                if os.path.exists("/usr/bin/bash"):  # WSL
-                    command = ["bash"] + command
-                else:
-                    # Use cmd.exe with bash while preserving all arguments
-                    command = ["cmd.exe", "/c", "bash"] + command
-
-        process = subprocess.Popen(
-            command,
-            shell=True,
-            universal_newlines=True,
-        )
-
-        process.wait()
-
     except Exception as e:
-        print(f"Error: {str(e)}")
+        log.exception(f"Unexpected error: {e}")
         return 1
     finally:
-        if process and process.poll() is None:
+        # Restore the original signal handler
-            process.terminate()
+        signal.signal(signal.SIGINT, original_sigint)
-            process.wait()
-    return process.returncode
-
-
-def run_command(command):
-    try:
-        result = subprocess.run(command, capture_output=True, text=True, check=True)
-        print("Script Output\n", result.stdout)
-        return result.returncode
-    except subprocess.CalledProcessError as e:
-        print("Error running script:", e)
-        print("Error output:", e.stderr)
-        return e.returncode

llama_stack/providers/inline/agents/meta_reference/agent_instance.py

@@ -614,118 +614,133 @@ class ChatAgent(ShieldRunnerMixin):
                     logger.debug(f"completion message with EOM (iter: {n_iter}): {str(message)}")
                     input_messages = input_messages + [message]
             else:
-                logger.debug(f"completion message (iter: {n_iter}) from the model: {str(message)}")
+                input_messages = input_messages + [message]
-                # 1. Start the tool execution step and progress
-                step_id = str(uuid.uuid4())
-                yield AgentTurnResponseStreamChunk(
-                    event=AgentTurnResponseEvent(
-                        payload=AgentTurnResponseStepStartPayload(
-                            step_type=StepType.tool_execution.value,
-                            step_id=step_id,
-                        )
-                    )
-                )
-                tool_call = message.tool_calls[0]
-                yield AgentTurnResponseStreamChunk(
-                    event=AgentTurnResponseEvent(
-                        payload=AgentTurnResponseStepProgressPayload(
-                            step_type=StepType.tool_execution.value,
-                            step_id=step_id,
-                            tool_call=tool_call,
-                            delta=ToolCallDelta(
-                                parse_status=ToolCallParseStatus.in_progress,
-                                tool_call=tool_call,
-                            ),
-                        )
-                    )
-                )
 
-                # If tool is a client tool, yield CompletionMessage and return
+                # Process tool calls in the message
-                if tool_call.tool_name in client_tools:
+                client_tool_calls = []
-                    # NOTE: mark end_of_message to indicate to client that it may
+                non_client_tool_calls = []
-                    # call the tool and continue the conversation with the tool's response.
+
-                    message.stop_reason = StopReason.end_of_message
+                # Separate client and non-client tool calls
+                for tool_call in message.tool_calls:
+                    if tool_call.tool_name in client_tools:
+                        client_tool_calls.append(tool_call)
+                    else:
+                        non_client_tool_calls.append(tool_call)
+
+                # Process non-client tool calls first
+                for tool_call in non_client_tool_calls:
+                    step_id = str(uuid.uuid4())
+                    yield AgentTurnResponseStreamChunk(
+                        event=AgentTurnResponseEvent(
+                            payload=AgentTurnResponseStepStartPayload(
+                                step_type=StepType.tool_execution.value,
+                                step_id=step_id,
+                            )
+                        )
+                    )
+
+                    yield AgentTurnResponseStreamChunk(
+                        event=AgentTurnResponseEvent(
+                            payload=AgentTurnResponseStepProgressPayload(
+                                step_type=StepType.tool_execution.value,
+                                step_id=step_id,
+                                delta=ToolCallDelta(
+                                    parse_status=ToolCallParseStatus.in_progress,
+                                    tool_call=tool_call,
+                                ),
+                            )
+                        )
+                    )
+
+                    # Execute the tool call
+                    async with tracing.span(
+                        "tool_execution",
+                        {
+                            "tool_name": tool_call.tool_name,
+                            "input": message.model_dump_json(),
+                        },
+                    ) as span:
+                        tool_execution_start_time = datetime.now(timezone.utc).isoformat()
+                        tool_result = await self.execute_tool_call_maybe(
+                            session_id,
+                            tool_call,
+                        )
+                        if tool_result.content is None:
+                            raise ValueError(
+                                f"Tool call result (id: {tool_call.call_id}, name: {tool_call.tool_name}) does not have any content"
+                            )
+                        result_message = ToolResponseMessage(
+                            call_id=tool_call.call_id,
+                            content=tool_result.content,
+                        )
+                        span.set_attribute("output", result_message.model_dump_json())
+
+                        # Store tool execution step
+                        tool_execution_step = ToolExecutionStep(
+                            step_id=step_id,
+                            turn_id=turn_id,
+                            tool_calls=[tool_call],
+                            tool_responses=[
+                                ToolResponse(
+                                    call_id=tool_call.call_id,
+                                    tool_name=tool_call.tool_name,
+                                    content=tool_result.content,
+                                    metadata=tool_result.metadata,
+                                )
+                            ],
+                            started_at=tool_execution_start_time,
+                            completed_at=datetime.now(timezone.utc).isoformat(),
+                        )
+
+                        # Yield the step completion event
+                        yield AgentTurnResponseStreamChunk(
+                            event=AgentTurnResponseEvent(
+                                payload=AgentTurnResponseStepCompletePayload(
+                                    step_type=StepType.tool_execution.value,
+                                    step_id=step_id,
+                                    step_details=tool_execution_step,
+                                )
+                            )
+                        )
+
+                        # Add the result message to input_messages for the next iteration
+                        input_messages.append(result_message)
+
+                        # TODO: add tool-input touchpoint and a "start" event for this step also
+                        # but that needs a lot more refactoring of Tool code potentially
+                        if (type(result_message.content) is str) and (
+                            out_attachment := _interpret_content_as_attachment(result_message.content)
+                        ):
+                            # NOTE: when we push this message back to the model, the model may ignore the
+                            # attached file path etc. since the model is trained to only provide a user message
+                            # with the summary. We keep all generated attachments and then attach them to final message
+                            output_attachments.append(out_attachment)
+
+                # If there are client tool calls, yield a message with only those tool calls
+                if client_tool_calls:
                     await self.storage.set_in_progress_tool_call_step(
                         session_id,
                         turn_id,
                         ToolExecutionStep(
                             step_id=step_id,
                             turn_id=turn_id,
-                            tool_calls=[tool_call],
+                            tool_calls=client_tool_calls,
                             tool_responses=[],
                             started_at=datetime.now(timezone.utc).isoformat(),
                         ),
                     )
-                    yield message
+
+                    # Create a copy of the message with only client tool calls
+                    client_message = message.model_copy(deep=True)
+                    client_message.tool_calls = client_tool_calls
+                    # NOTE: mark end_of_message to indicate to client that it may
+                    # call the tool and continue the conversation with the tool's response.
+                    client_message.stop_reason = StopReason.end_of_message
+
+                    # Yield the message with client tool calls
+                    yield client_message
                     return
 
-                # If tool is a builtin server tool, execute it
-                tool_name = tool_call.tool_name
-                if isinstance(tool_name, BuiltinTool):
-                    tool_name = tool_name.value
-                async with tracing.span(
-                    "tool_execution",
-                    {
-                        "tool_name": tool_name,
-                        "input": message.model_dump_json(),
-                    },
-                ) as span:
-                    tool_execution_start_time = datetime.now(timezone.utc).isoformat()
-                    tool_call = message.tool_calls[0]
-                    tool_result = await self.execute_tool_call_maybe(
-                        session_id,
-                        tool_call,
-                    )
-                    if tool_result.content is None:
-                        raise ValueError(
-                            f"Tool call result (id: {tool_call.call_id}, name: {tool_call.tool_name}) does not have any content"
-                        )
-                    result_messages = [
-                        ToolResponseMessage(
-                            call_id=tool_call.call_id,
-                            content=tool_result.content,
-                        )
-                    ]
-                    assert len(result_messages) == 1, "Currently not supporting multiple messages"
-                    result_message = result_messages[0]
-                    span.set_attribute("output", result_message.model_dump_json())
-
-                yield AgentTurnResponseStreamChunk(
-                    event=AgentTurnResponseEvent(
-                        payload=AgentTurnResponseStepCompletePayload(
-                            step_type=StepType.tool_execution.value,
-                            step_id=step_id,
-                            step_details=ToolExecutionStep(
-                                step_id=step_id,
-                                turn_id=turn_id,
-                                tool_calls=[tool_call],
-                                tool_responses=[
-                                    ToolResponse(
-                                        call_id=result_message.call_id,
-                                        tool_name=tool_call.tool_name,
-                                        content=result_message.content,
-                                        metadata=tool_result.metadata,
-                                    )
-                                ],
-                                started_at=tool_execution_start_time,
-                                completed_at=datetime.now(timezone.utc).isoformat(),
-                            ),
-                        )
-                    )
-                )
-
-                # TODO: add tool-input touchpoint and a "start" event for this step also
-                # but that needs a lot more refactoring of Tool code potentially
-                if (type(result_message.content) is str) and (
-                    out_attachment := _interpret_content_as_attachment(result_message.content)
-                ):
-                    # NOTE: when we push this message back to the model, the model may ignore the
-                    # attached file path etc. since the model is trained to only provide a user message
-                    # with the summary. We keep all generated attachments and then attach them to final message
-                    output_attachments.append(out_attachment)
-
-                input_messages = input_messages + [message, result_message]
-
     async def _initialize_tools(
         self,
         toolgroups_for_turn: Optional[List[AgentToolGroup]] = None,
@@ -891,16 +906,14 @@ class ChatAgent(ShieldRunnerMixin):
         if memory_tool and code_interpreter_tool:
             # if both memory and code_interpreter are available, we download the URLs
             # and attach the data to the last message.
-            msg = await attachment_message(self.tempdir, url_items)
+            await attachment_message(self.tempdir, url_items, input_messages[-1])
-            input_messages.append(msg)
             # Since memory is present, add all the data to the memory bank
             await self.add_to_session_vector_db(session_id, documents)
         elif code_interpreter_tool:
             # if only code_interpreter is available, we download the URLs to a tempdir
             # and attach the path to them as a message to inference with the
             # assumption that the model invokes the code_interpreter tool with the path
-            msg = await attachment_message(self.tempdir, url_items)
+            await attachment_message(self.tempdir, url_items, input_messages[-1])
-            input_messages.append(msg)
         elif memory_tool:
             # if only memory is available, we load the data from the URLs and content items to the memory bank
             await self.add_to_session_vector_db(session_id, documents)
@@ -967,8 +980,8 @@ async def load_data_from_urls(urls: List[URL]) -> List[str]:
     return data
 
 
-async def attachment_message(tempdir: str, urls: List[URL]) -> ToolResponseMessage:
+async def attachment_message(tempdir: str, urls: List[URL], message: UserMessage) -> None:
-    content = []
+    contents = []
 
     for url in urls:
         uri = url.uri
@@ -988,16 +1001,19 @@ async def attachment_message(tempdir: str, urls: List[URL]) -> ToolResponseMessa
         else:
             raise ValueError(f"Unsupported URL {url}")
 
-        content.append(
+        contents.append(
             TextContentItem(
                 text=f'# User provided a file accessible to you at "{filepath}"\nYou can use code_interpreter to load and inspect it.'
             )
         )
 
-    return ToolResponseMessage(
+    if isinstance(message.content, list):
-        call_id="",
+        message.content.extend(contents)
-        content=content,
+    else:
-    )
+        if isinstance(message.content, str):
+            message.content = [TextContentItem(text=message.content)] + contents
+        else:
+            message.content = [message.content] + contents
 
 
 def _interpret_content_as_attachment(

llama_stack/providers/inline/datasetio/localfs/datasetio.py

@@ -3,20 +3,14 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-import base64
-import os
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
-from urllib.parse import urlparse
 
 import pandas
 
-from llama_stack.apis.common.content_types import URL
+from llama_stack.apis.datasetio import DatasetIO, IterrowsResponse
-from llama_stack.apis.datasetio import DatasetIO, PaginatedRowsResult
 from llama_stack.apis.datasets import Dataset
 from llama_stack.providers.datatypes import DatasetsProtocolPrivate
-from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_url
+from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_uri
 from llama_stack.providers.utils.kvstore import kvstore_impl
 
 from .config import LocalFSDatasetIOConfig
@@ -24,30 +18,7 @@ from .config import LocalFSDatasetIOConfig
 DATASETS_PREFIX = "localfs_datasets:"
 
 
-class BaseDataset(ABC):
+class PandasDataframeDataset:
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
-
-    @abstractmethod
-    def __len__(self) -> int:
-        raise NotImplementedError()
-
-    @abstractmethod
-    def __getitem__(self, idx):
-        raise NotImplementedError()
-
-    @abstractmethod
-    def load(self):
-        raise NotImplementedError()
-
-
-@dataclass
-class DatasetInfo:
-    dataset_def: Dataset
-    dataset_impl: BaseDataset
-
-
-class PandasDataframeDataset(BaseDataset):
     def __init__(self, dataset_def: Dataset, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.dataset_def = dataset_def
@@ -64,23 +35,19 @@ class PandasDataframeDataset(BaseDataset):
         else:
             return self.df.iloc[idx].to_dict()
 
-    def _validate_dataset_schema(self, df) -> pandas.DataFrame:
-        # note that we will drop any columns in dataset that are not in the schema
-        df = df[self.dataset_def.dataset_schema.keys()]
-        # check all columns in dataset schema are present
-        assert len(df.columns) == len(self.dataset_def.dataset_schema)
-        # TODO: type checking against column types in dataset schema
-        return df
-
     def load(self) -> None:
         if self.df is not None:
             return
 
-        df = get_dataframe_from_url(self.dataset_def.url)
+        if self.dataset_def.source.type == "uri":
-        if df is None:
+            self.df = get_dataframe_from_uri(self.dataset_def.source.uri)
-            raise ValueError(f"Failed to load dataset from {self.dataset_def.url}")
+        elif self.dataset_def.source.type == "rows":
+            self.df = pandas.DataFrame(self.dataset_def.source.rows)
+        else:
+            raise ValueError(f"Unsupported dataset source type: {self.dataset_def.source.type}")
 
-        self.df = self._validate_dataset_schema(df)
+        if self.df is None:
+            raise ValueError(f"Failed to load dataset from {self.dataset_def.url}")
 
 
 class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
@@ -99,95 +66,55 @@ class LocalFSDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
 
         for dataset in stored_datasets:
             dataset = Dataset.model_validate_json(dataset)
-            dataset_impl = PandasDataframeDataset(dataset)
+            self.dataset_infos[dataset.identifier] = dataset
-            self.dataset_infos[dataset.identifier] = DatasetInfo(
-                dataset_def=dataset,
-                dataset_impl=dataset_impl,
-            )
 
     async def shutdown(self) -> None: ...
 
     async def register_dataset(
         self,
-        dataset: Dataset,
+        dataset_def: Dataset,
     ) -> None:
         # Store in kvstore
-        key = f"{DATASETS_PREFIX}{dataset.identifier}"
+        key = f"{DATASETS_PREFIX}{dataset_def.identifier}"
         await self.kvstore.set(
             key=key,
-            value=dataset.json(),
+            value=dataset_def.model_dump_json(),
-        )
-        dataset_impl = PandasDataframeDataset(dataset)
-        self.dataset_infos[dataset.identifier] = DatasetInfo(
-            dataset_def=dataset,
-            dataset_impl=dataset_impl,
         )
+        self.dataset_infos[dataset_def.identifier] = dataset_def
 
     async def unregister_dataset(self, dataset_id: str) -> None:
         key = f"{DATASETS_PREFIX}{dataset_id}"
         await self.kvstore.delete(key=key)
         del self.dataset_infos[dataset_id]
 
-    async def get_rows_paginated(
+    async def iterrows(
         self,
         dataset_id: str,
-        rows_in_page: int,
+        start_index: Optional[int] = None,
-        page_token: Optional[str] = None,
+        limit: Optional[int] = None,
-        filter_condition: Optional[str] = None,
+    ) -> IterrowsResponse:
-    ) -> PaginatedRowsResult:
+        dataset_def = self.dataset_infos[dataset_id]
-        dataset_info = self.dataset_infos.get(dataset_id)
+        dataset_impl = PandasDataframeDataset(dataset_def)
-        dataset_info.dataset_impl.load()
+        dataset_impl.load()
 
-        if page_token and not page_token.isnumeric():
+        start_index = start_index or 0
-            raise ValueError("Invalid page_token")
 
-        if page_token is None or len(page_token) == 0:
+        if limit is None or limit == -1:
-            next_page_token = 0
+            end = len(dataset_impl)
         else:
-            next_page_token = int(page_token)
+            end = min(start_index + limit, len(dataset_impl))
 
-        start = next_page_token
+        rows = dataset_impl[start_index:end]
-        if rows_in_page == -1:
-            end = len(dataset_info.dataset_impl)
-        else:
-            end = min(start + rows_in_page, len(dataset_info.dataset_impl))
 
-        rows = dataset_info.dataset_impl[start:end]
+        return IterrowsResponse(
-
+            data=rows,
-        return PaginatedRowsResult(
+            next_start_index=end if end < len(dataset_impl) else None,
-            rows=rows,
-            total_count=len(rows),
-            next_page_token=str(end),
         )
 
     async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:
-        dataset_info = self.dataset_infos.get(dataset_id)
+        dataset_def = self.dataset_infos[dataset_id]
-        if dataset_info is None:
+        dataset_impl = PandasDataframeDataset(dataset_def)
-            raise ValueError(f"Dataset with id {dataset_id} not found")
-
-        dataset_impl = dataset_info.dataset_impl
         dataset_impl.load()
 
         new_rows_df = pandas.DataFrame(rows)
-        new_rows_df = dataset_impl._validate_dataset_schema(new_rows_df)
         dataset_impl.df = pandas.concat([dataset_impl.df, new_rows_df], ignore_index=True)
-
-        url = str(dataset_info.dataset_def.url.uri)
-        parsed_url = urlparse(url)
-
-        if parsed_url.scheme == "file" or not parsed_url.scheme:
-            file_path = parsed_url.path
-            os.makedirs(os.path.dirname(file_path), exist_ok=True)
-            dataset_impl.df.to_csv(file_path, index=False)
-        elif parsed_url.scheme == "data":
-            # For data URLs, we need to update the base64-encoded content
-            if not parsed_url.path.startswith("text/csv;base64,"):
-                raise ValueError("Data URL must be a base64-encoded CSV")
-
-            csv_buffer = dataset_impl.df.to_csv(index=False)
-            base64_content = base64.b64encode(csv_buffer.encode("utf-8")).decode("utf-8")
-            dataset_info.dataset_def.url = URL(uri=f"data:text/csv;base64,{base64_content}")
-        else:
-            raise ValueError(
-                f"Unsupported URL scheme: {parsed_url.scheme}. Only file:// and data: URLs are supported for writing."
-            )

llama_stack/providers/inline/eval/meta_reference/eval.py

@@ -14,16 +14,11 @@ from llama_stack.apis.datasetio import DatasetIO
 from llama_stack.apis.datasets import Datasets
 from llama_stack.apis.inference import Inference, SystemMessage, UserMessage
 from llama_stack.apis.scoring import Scoring
-from llama_stack.distribution.datatypes import Api
 from llama_stack.providers.datatypes import BenchmarksProtocolPrivate
 from llama_stack.providers.inline.agents.meta_reference.agent_instance import (
     MEMORY_QUERY_TOOL,
 )
-from llama_stack.providers.utils.common.data_schema_validator import (
+from llama_stack.providers.utils.common.data_schema_validator import ColumnName
-    ColumnName,
-    get_valid_schemas,
-    validate_dataset_schema,
-)
 from llama_stack.providers.utils.kvstore import kvstore_impl
 
 from .....apis.common.job_types import Job
@@ -88,15 +83,17 @@ class MetaReferenceEvalImpl(
         task_def = self.benchmarks[benchmark_id]
         dataset_id = task_def.dataset_id
         scoring_functions = task_def.scoring_functions
-        dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
+
-        validate_dataset_schema(dataset_def.dataset_schema, get_valid_schemas(Api.eval.value))
+        # TODO (xiyan): validate dataset schema
-        all_rows = await self.datasetio_api.get_rows_paginated(
+        # dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
+
+        all_rows = await self.datasetio_api.iterrows(
             dataset_id=dataset_id,
-            rows_in_page=(-1 if benchmark_config.num_examples is None else benchmark_config.num_examples),
+            limit=(-1 if benchmark_config.num_examples is None else benchmark_config.num_examples),
         )
         res = await self.evaluate_rows(
             benchmark_id=benchmark_id,
-            input_rows=all_rows.rows,
+            input_rows=all_rows.data,
             scoring_functions=scoring_functions,
             benchmark_config=benchmark_config,
         )

llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py

@@ -328,13 +328,13 @@ class LoraFinetuningSingleDevice:
         batch_size: int,
     ) -> Tuple[DistributedSampler, DataLoader]:
         async def fetch_rows(dataset_id: str):
-            return await self.datasetio_api.get_rows_paginated(
+            return await self.datasetio_api.iterrows(
                 dataset_id=dataset_id,
-                rows_in_page=-1,
+                limit=-1,
             )
 
         all_rows = await fetch_rows(dataset_id)
-        rows = all_rows.rows
+        rows = all_rows.data
 
         await validate_input_dataset_schema(
             datasets_api=self.datasets_api,

llama_stack/providers/inline/safety/llama_guard/llama_guard.py

@@ -227,13 +227,6 @@ class LlamaGuardShield:
         if len(messages) >= 2 and (messages[0].role == Role.user.value and messages[1].role == Role.user.value):
             messages = messages[1:]
 
-        for i in range(1, len(messages)):
-            if messages[i].role == messages[i - 1].role:
-                for i, m in enumerate(messages):
-                    print(f"{i}: {m.role}: {m.content}")
-                raise ValueError(
-                    f"Messages must alternate between user and assistant. Message {i} has the same role as message {i - 1}"
-                )
         return messages
 
     async def run(self, messages: List[Message]) -> RunShieldResponse:

llama_stack/providers/inline/scoring/basic/scoring.py

@@ -24,7 +24,9 @@ from llama_stack.providers.utils.common.data_schema_validator import (
 from .config import BasicScoringConfig
 from .scoring_fn.bfcl_scoring_fn import BFCLScoringFn
 from .scoring_fn.equality_scoring_fn import EqualityScoringFn
-from .scoring_fn.regex_parser_math_response_scoring_fn import RegexParserMathResponseScoringFn
+from .scoring_fn.regex_parser_math_response_scoring_fn import (
+    RegexParserMathResponseScoringFn,
+)
 from .scoring_fn.regex_parser_scoring_fn import RegexParserScoringFn
 from .scoring_fn.subset_of_scoring_fn import SubsetOfScoringFn
 
@@ -82,12 +84,12 @@ class BasicScoringImpl(
         dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
         validate_dataset_schema(dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value))
 
-        all_rows = await self.datasetio_api.get_rows_paginated(
+        all_rows = await self.datasetio_api.iterrows(
             dataset_id=dataset_id,
-            rows_in_page=-1,
+            limit=-1,
         )
         res = await self.score(
-            input_rows=all_rows.rows,
+            input_rows=all_rows.data,
             scoring_functions=scoring_functions,
         )
         if save_results_dataset:

llama_stack/providers/inline/scoring/braintrust/braintrust.py

@@ -167,11 +167,11 @@ class BraintrustScoringImpl(
         dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
         validate_dataset_schema(dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value))
 
-        all_rows = await self.datasetio_api.get_rows_paginated(
+        all_rows = await self.datasetio_api.iterrows(
             dataset_id=dataset_id,
-            rows_in_page=-1,
+            limit=-1,
         )
-        res = await self.score(input_rows=all_rows.rows, scoring_functions=scoring_functions)
+        res = await self.score(input_rows=all_rows.data, scoring_functions=scoring_functions)
         if save_results_dataset:
             # TODO: persist and register dataset on to server for reading
             # self.datasets_api.register_dataset()

llama_stack/providers/inline/scoring/llm_as_judge/scoring.py

@@ -72,12 +72,12 @@ class LlmAsJudgeScoringImpl(
         dataset_def = await self.datasets_api.get_dataset(dataset_id=dataset_id)
         validate_dataset_schema(dataset_def.dataset_schema, get_valid_schemas(Api.scoring.value))
 
-        all_rows = await self.datasetio_api.get_rows_paginated(
+        all_rows = await self.datasetio_api.iterrows(
             dataset_id=dataset_id,
-            rows_in_page=-1,
+            limit=-1,
         )
         res = await self.score(
-            input_rows=all_rows.rows,
+            input_rows=all_rows.data,
             scoring_functions=scoring_functions,
         )
         if save_results_dataset:

llama_stack/providers/registry/safety.py

@@ -55,4 +55,13 @@ def available_providers() -> List[ProviderSpec]:
                 config_class="llama_stack.providers.remote.safety.bedrock.BedrockSafetyConfig",
             ),
         ),
+        remote_provider_spec(
+            api=Api.safety,
+            adapter=AdapterSpec(
+                adapter_type="nvidia",
+                pip_packages=["requests"],
+                module="llama_stack.providers.remote.safety.nvidia",
+                config_class="llama_stack.providers.remote.safety.nvidia.NVIDIASafetyConfig",
+            ),
+        ),
     ]

llama_stack/providers/remote/datasetio/huggingface/huggingface.py

@@ -4,13 +4,13 @@
 # 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, Dict, List, Optional
+from urllib.parse import parse_qs, urlparse
 
 import datasets as hf_datasets
 
-from llama_stack.apis.datasetio import DatasetIO, PaginatedRowsResult
+from llama_stack.apis.datasetio import DatasetIO, IterrowsResponse
 from llama_stack.apis.datasets import Dataset
 from llama_stack.providers.datatypes import DatasetsProtocolPrivate
-from llama_stack.providers.utils.datasetio.url_utils import get_dataframe_from_url
 from llama_stack.providers.utils.kvstore import kvstore_impl
 
 from .config import HuggingfaceDatasetIOConfig
@@ -18,22 +18,14 @@ from .config import HuggingfaceDatasetIOConfig
 DATASETS_PREFIX = "datasets:"
 
 
-def load_hf_dataset(dataset_def: Dataset):
+def parse_hf_params(dataset_def: Dataset):
-    if dataset_def.metadata.get("path", None):
+    uri = dataset_def.source.uri
-        dataset = hf_datasets.load_dataset(**dataset_def.metadata)
+    parsed_uri = urlparse(uri)
-    else:
+    params = parse_qs(parsed_uri.query)
-        df = get_dataframe_from_url(dataset_def.url)
+    params = {k: v[0] for k, v in params.items()}
+    path = parsed_uri.path.lstrip("/")
 
-        if df is None:
+    return path, params
-            raise ValueError(f"Failed to load dataset from {dataset_def.url}")
-
-        dataset = hf_datasets.Dataset.from_pandas(df)
-
-    # drop columns not specified by schema
-    if dataset_def.dataset_schema:
-        dataset = dataset.select_columns(list(dataset_def.dataset_schema.keys()))
-
-    return dataset
 
 
 class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
@@ -64,7 +56,7 @@ class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
         key = f"{DATASETS_PREFIX}{dataset_def.identifier}"
         await self.kvstore.set(
             key=key,
-            value=dataset_def.json(),
+            value=dataset_def.model_dump_json(),
         )
         self.dataset_infos[dataset_def.identifier] = dataset_def
 
@@ -73,41 +65,34 @@ class HuggingfaceDatasetIOImpl(DatasetIO, DatasetsProtocolPrivate):
         await self.kvstore.delete(key=key)
         del self.dataset_infos[dataset_id]
 
-    async def get_rows_paginated(
+    async def iterrows(
         self,
         dataset_id: str,
-        rows_in_page: int,
+        start_index: Optional[int] = None,
-        page_token: Optional[str] = None,
+        limit: Optional[int] = None,
-        filter_condition: Optional[str] = None,
+    ) -> IterrowsResponse:
-    ) -> PaginatedRowsResult:
         dataset_def = self.dataset_infos[dataset_id]
-        loaded_dataset = load_hf_dataset(dataset_def)
+        path, params = parse_hf_params(dataset_def)
+        loaded_dataset = hf_datasets.load_dataset(path, **params)
 
-        if page_token and not page_token.isnumeric():
+        start_index = start_index or 0
-            raise ValueError("Invalid page_token")
 
-        if page_token is None or len(page_token) == 0:
+        if limit is None or limit == -1:
-            next_page_token = 0
-        else:
-            next_page_token = int(page_token)
-
-        start = next_page_token
-        if rows_in_page == -1:
             end = len(loaded_dataset)
         else:
-            end = min(start + rows_in_page, len(loaded_dataset))
+            end = min(start_index + limit, len(loaded_dataset))
 
-        rows = [loaded_dataset[i] for i in range(start, end)]
+        rows = [loaded_dataset[i] for i in range(start_index, end)]
 
-        return PaginatedRowsResult(
+        return IterrowsResponse(
-            rows=rows,
+            data=rows,
-            total_count=len(rows),
+            next_start_index=end if end < len(loaded_dataset) else None,
-            next_page_token=str(end),
         )
 
     async def append_rows(self, dataset_id: str, rows: List[Dict[str, Any]]) -> None:
         dataset_def = self.dataset_infos[dataset_id]
-        loaded_dataset = load_hf_dataset(dataset_def)
+        path, params = parse_hf_params(dataset_def)
+        loaded_dataset = hf_datasets.load_dataset(path, **params)
 
         # Convert rows to HF Dataset format
         new_dataset = hf_datasets.Dataset.from_list(rows)

llama_stack/providers/remote/inference/passthrough/passthrough.py

@@ -12,6 +12,7 @@ from llama_stack.apis.common.content_types import InterleavedContent
 from llama_stack.apis.inference import (
     ChatCompletionResponse,
     ChatCompletionResponseStreamChunk,
+    CompletionMessage,
     EmbeddingsResponse,
     EmbeddingTaskType,
     Inference,
@@ -160,12 +161,14 @@ class PassthroughInferenceAdapter(Inference):
         client = self._get_client()
         response = await client.inference.chat_completion(**json_params)
 
-        response = response.to_dict()
+        return ChatCompletionResponse(
-
+            completion_message=CompletionMessage(
-        # temporary hack to remove the metrics from the response
+                content=response.completion_message.content.text,
-        response["metrics"] = []
+                stop_reason=response.completion_message.stop_reason,
-
+                tool_calls=response.completion_message.tool_calls,
-        return convert_to_pydantic(ChatCompletionResponse, response)
+            ),
+            logprobs=response.logprobs,
+        )
 
     async def _stream_chat_completion(self, json_params: Dict[str, Any]) -> AsyncGenerator:
         client = self._get_client()

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

@@ -25,6 +25,10 @@ class VLLMInferenceAdapterConfig(BaseModel):
         default="fake",
         description="The API token",
     )
+    tls_verify: bool = Field(
+        default=True,
+        description="Whether to verify TLS certificates",
+    )
 
     @classmethod
     def sample_run_config(
@@ -36,4 +40,5 @@ class VLLMInferenceAdapterConfig(BaseModel):
             "url": 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,6 +7,7 @@ import json
 import logging
 from typing import AsyncGenerator, List, Optional, Union
 
+import httpx
 from openai import AsyncOpenAI
 from openai.types.chat.chat_completion_chunk import (
     ChatCompletionChunk as OpenAIChatCompletionChunk,
@@ -229,7 +230,11 @@ class VLLMInferenceAdapter(Inference, ModelsProtocolPrivate):
 
     async def initialize(self) -> None:
         log.info(f"Initializing VLLM client with base_url={self.config.url}")
-        self.client = AsyncOpenAI(base_url=self.config.url, api_key=self.config.api_token)
+        self.client = AsyncOpenAI(
+            base_url=self.config.url,
+            api_key=self.config.api_token,
+            http_client=None if self.config.tls_verify else httpx.AsyncClient(verify=False),
+        )
 
     async def shutdown(self) -> None:
         pass

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

@@ -0,0 +1,18 @@
+# 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 .config import NVIDIASafetyConfig
+
+
+async def get_adapter_impl(config: NVIDIASafetyConfig, _deps) -> Any:
+    from .nvidia import NVIDIASafetyAdapter
+
+    impl = NVIDIASafetyAdapter(config)
+    await impl.initialize()
+    return impl

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

@@ -0,0 +1,37 @@
+# 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, Dict, Optional
+
+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: Optional[str] = Field(default="self-check", description="Config ID to use from the config store")
+
+    @classmethod
+    def sample_run_config(cls, **kwargs) -> Dict[str, Any]:
+        return {
+            "guardrails_service_url": "${env.GUARDRAILS_SERVICE_URL:http://localhost:7331}",
+            "config_id": "self-check",
+        }

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

@@ -0,0 +1,154 @@
+# 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 logging
+from typing import Any, List, Optional
+
+import requests
+
+from llama_stack.apis.inference import Message
+from llama_stack.apis.safety import RunShieldResponse, Safety, SafetyViolation, ViolationLevel
+from llama_stack.apis.shields import Shield
+from llama_stack.distribution.library_client import convert_pydantic_to_json_value
+from llama_stack.providers.datatypes import ShieldsProtocolPrivate
+
+from .config import NVIDIASafetyConfig
+
+logger = logging.getLogger(__name__)
+
+
+class NVIDIASafetyAdapter(Safety, ShieldsProtocolPrivate):
+    def __init__(self, config: NVIDIASafetyConfig) -> None:
+        """
+        Initialize the NVIDIASafetyAdapter with a given safety configuration.
+
+        Args:
+            config (NVIDIASafetyConfig): The configuration containing the guardrails service URL and config ID.
+        """
+        print(f"Initializing NVIDIASafetyAdapter({config.guardrails_service_url})...")
+        self.config = config
+
+    async def initialize(self) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        pass
+
+    async def register_shield(self, shield: Shield) -> None:
+        if not shield.provider_resource_id:
+            raise ValueError("Shield model not provided.")
+
+    async def run_shield(
+        self, shield_id: str, messages: List[Message], params: Optional[dict[str, Any]] = 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)
+        if not shield:
+            raise ValueError(f"Shield {shield_id} not found")
+
+        self.shield = NeMoGuardrails(self.config, shield.shield_id)
+        return await self.shield.run(messages)
+
+
+class NeMoGuardrails:
+    """
+    A class that encapsulates NVIDIA's guardrails safety logic.
+
+    Sends messages to the guardrails service and interprets the response to determine
+    if a safety violation has occurred.
+    """
+
+    def __init__(
+        self,
+        config: NVIDIASafetyConfig,
+        model: str,
+        threshold: float = 0.9,
+        temperature: float = 1.0,
+    ):
+        """
+        Initialize a NeMoGuardrails instance with the provided parameters.
+
+        Args:
+            config (NVIDIASafetyConfig): The safety configuration containing the config ID and guardrails URL.
+            model (str): The identifier or name of the model to be used for safety checks.
+            threshold (float, optional): The threshold for flagging violations. Defaults to 0.9.
+            temperature (float, optional): The temperature setting for the underlying model. Must be greater than 0. Defaults to 1.0.
+
+        Raises:
+            ValueError: If temperature is less than or equal to 0.
+            AssertionError: If config_id is not provided in the configuration.
+        """
+        self.config_id = config.config_id
+        self.model = model
+        assert self.config_id is not None, "Must provide config id"
+        if temperature <= 0:
+            raise ValueError("Temperature must be greater than 0")
+
+        self.temperature = temperature
+        self.threshold = threshold
+        self.guardrails_service_url = config.guardrails_service_url
+
+    async def run(self, messages: List[Message]) -> RunShieldResponse:
+        """
+        Queries the /v1/guardrails/checks endpoint of the NeMo guardrails deployed API.
+
+        Args:
+            messages (List[Message]): A list of Message objects to be checked for safety violations.
+
+        Returns:
+            RunShieldResponse: If the response indicates a violation ("blocked" status), returns a
+            RunShieldResponse with a SafetyViolation; otherwise, returns a RunShieldResponse with violation set to None.
+
+        Raises:
+            requests.HTTPError: If the POST request fails.
+        """
+        headers = {
+            "Accept": "application/json",
+        }
+        request_data = {
+            "model": self.model,
+            "messages": convert_pydantic_to_json_value(messages),
+            "temperature": self.temperature,
+            "top_p": 1,
+            "frequency_penalty": 0,
+            "presence_penalty": 0,
+            "max_tokens": 160,
+            "stream": False,
+            "guardrails": {
+                "config_id": self.config_id,
+            },
+        }
+        response = requests.post(
+            url=f"{self.guardrails_service_url}/v1/guardrail/checks", headers=headers, json=request_data
+        )
+        response.raise_for_status()
+        if "Content-Type" in response.headers and response.headers["Content-Type"].startswith("application/json"):
+            response_json = response.json()
+        if response_json["status"] == "blocked":
+            user_message = "Sorry I cannot do this."
+            metadata = response_json["rails_status"]
+
+            return RunShieldResponse(
+                violation=SafetyViolation(
+                    user_message=user_message,
+                    violation_level=ViolationLevel.ERROR,
+                    metadata=metadata,
+                )
+            )
+        return RunShieldResponse(violation=None)

llama_stack/providers/utils/datasetio/url_utils.py

@@ -10,18 +10,17 @@ from urllib.parse import unquote
 
 import pandas
 
-from llama_stack.apis.common.content_types import URL
 from llama_stack.providers.utils.memory.vector_store import parse_data_url
 
 
-def get_dataframe_from_url(url: URL):
+def get_dataframe_from_uri(uri: str):
     df = None
-    if url.uri.endswith(".csv"):
+    if uri.endswith(".csv"):
-        df = pandas.read_csv(url.uri)
+        df = pandas.read_csv(uri)
-    elif url.uri.endswith(".xlsx"):
+    elif uri.endswith(".xlsx"):
-        df = pandas.read_excel(url.uri)
+        df = pandas.read_excel(uri)
-    elif url.uri.startswith("data:"):
+    elif uri.startswith("data:"):
-        parts = parse_data_url(url.uri)
+        parts = parse_data_url(uri)
         data = parts["data"]
         if parts["is_base64"]:
             data = base64.b64decode(data)
@@ -39,6 +38,6 @@ def get_dataframe_from_url(url: URL):
         else:
             df = pandas.read_excel(data_bytes)
     else:
-        raise ValueError(f"Unsupported file type: {url}")
+        raise ValueError(f"Unsupported file type: {uri}")
 
     return df

llama_stack/providers/utils/inference/litellm_openai_mixin.py

@@ -192,7 +192,11 @@ class LiteLLMOpenAIMixin(
         if request.tools:
             input_dict["tools"] = [convert_tooldef_to_openai_tool(tool) for tool in request.tools]
             if request.tool_config.tool_choice:
-                input_dict["tool_choice"] = request.tool_config.tool_choice.value
+                input_dict["tool_choice"] = (
+                    request.tool_config.tool_choice.value
+                    if isinstance(request.tool_config.tool_choice, ToolChoice)
+                    else request.tool_config.tool_choice
+                )
 
         provider_data = self.get_request_provider_data()
         key_field = self.provider_data_api_key_field

llama_stack/providers/utils/inference/openai_compat.py

@@ -527,27 +527,31 @@ async def convert_message_to_openai_dict_new(
     async def _convert_message_content(
         content: InterleavedContent,
     ) -> Union[str, Iterable[OpenAIChatCompletionContentPartParam]]:
-        async def impl():
+        async def impl(
+            content_: InterleavedContent,
+        ) -> Union[str, OpenAIChatCompletionContentPartParam, List[OpenAIChatCompletionContentPartParam]]:
             # Llama Stack and OpenAI spec match for str and text input
-            if isinstance(content, str):
+            if isinstance(content_, str):
-                return content
+                return content_
-            elif isinstance(content, TextContentItem):
+            elif isinstance(content_, TextContentItem):
                 return OpenAIChatCompletionContentPartTextParam(
                     type="text",
-                    text=content.text,
+                    text=content_.text,
                 )
-            elif isinstance(content, ImageContentItem):
+            elif isinstance(content_, ImageContentItem):
                 return OpenAIChatCompletionContentPartImageParam(
                     type="image_url",
-                    image_url=OpenAIImageURL(url=await convert_image_content_to_url(content)),
+                    image_url=OpenAIImageURL(url=await convert_image_content_to_url(content_)),
                 )
-            elif isinstance(content, list):
+            elif isinstance(content_, list):
-                return [await _convert_message_content(item) for item in content]
+                return [await impl(item) for item in content_]
             else:
-                raise ValueError(f"Unsupported content type: {type(content)}")
+                raise ValueError(f"Unsupported content type: {type(content_)}")
 
-        ret = await impl()
+        ret = await impl(content)
-        if isinstance(ret, str) or isinstance(ret, list) or isinstance(ret, dict):
+
+        # OpenAI*Message expects a str or list
+        if isinstance(ret, str) or isinstance(ret, list):
             return ret
         else:
             return [ret]
@@ -566,13 +570,14 @@ async def convert_message_to_openai_dict_new(
                 OpenAIChatCompletionMessageToolCall(
                     id=tool.call_id,
                     function=OpenAIFunction(
-                        name=tool.tool_name,
+                        name=tool.tool_name if not isinstance(tool.tool_name, BuiltinTool) else tool.tool_name.value,
                         arguments=json.dumps(tool.arguments),
                     ),
                     type="function",
                 )
                 for tool in message.tool_calls
-            ],
+            ]
+            or None,
         )
     elif isinstance(message, ToolResponseMessage):
         out = OpenAIChatCompletionToolMessage(
@@ -858,7 +863,8 @@ async def convert_openai_chat_completion_stream(
     event_type = ChatCompletionResponseEventType.progress
 
     stop_reason = None
-    toolcall_buffer = {}
+    tool_call_idx_to_buffer = {}
+
     async for chunk in stream:
         choice = chunk.choices[0]  # assuming only one choice per chunk
 
@@ -868,7 +874,6 @@ async def convert_openai_chat_completion_stream(
 
         # if there's a tool call, emit an event for each tool in the list
         # if tool call and content, emit both separately
-
         if choice.delta.tool_calls:
             # the call may have content and a tool call. ChatCompletionResponseEvent
             # does not support both, so we emit the content first
@@ -889,44 +894,53 @@ async def convert_openai_chat_completion_stream(
                 )
 
             if not enable_incremental_tool_calls:
-                yield ChatCompletionResponseStreamChunk(
+                for tool_call in choice.delta.tool_calls:
-                    event=ChatCompletionResponseEvent(
+                    yield ChatCompletionResponseStreamChunk(
-                        event_type=next(event_type),
+                        event=ChatCompletionResponseEvent(
-                        delta=ToolCallDelta(
+                            event_type=event_type,
-                            tool_call=_convert_openai_tool_calls(choice.delta.tool_calls)[0],
+                            delta=ToolCallDelta(
-                            parse_status=ToolCallParseStatus.succeeded,
+                                tool_call=_convert_openai_tool_calls([tool_call])[0],
-                        ),
+                                parse_status=ToolCallParseStatus.succeeded,
-                        logprobs=_convert_openai_logprobs(logprobs),
+                            ),
+                            logprobs=_convert_openai_logprobs(logprobs),
+                        )
                     )
-                )
             else:
-                tool_call = choice.delta.tool_calls[0]
+                for tool_call in choice.delta.tool_calls:
-                if "name" not in toolcall_buffer:
+                    idx = tool_call.index if hasattr(tool_call, "index") else 0
-                    toolcall_buffer["call_id"] = tool_call.id
-                    toolcall_buffer["name"] = None
-                    toolcall_buffer["content"] = ""
-                if "arguments" not in toolcall_buffer:
-                    toolcall_buffer["arguments"] = ""
 
-                if tool_call.function.name:
+                    if idx not in tool_call_idx_to_buffer:
-                    toolcall_buffer["name"] = tool_call.function.name
+                        tool_call_idx_to_buffer[idx] = {
-                    delta = f"{toolcall_buffer['name']}("
+                            "call_id": tool_call.id,
-                if tool_call.function.arguments:
+                            "name": None,
-                    toolcall_buffer["arguments"] += tool_call.function.arguments
+                            "arguments": "",
-                    delta = toolcall_buffer["arguments"]
+                            "content": "",
+                        }
 
-                toolcall_buffer["content"] += delta
+                    buffer = tool_call_idx_to_buffer[idx]
-                yield ChatCompletionResponseStreamChunk(
+
-                    event=ChatCompletionResponseEvent(
+                    if tool_call.function:
-                        event_type=event_type,
+                        if tool_call.function.name:
-                        delta=ToolCallDelta(
+                            buffer["name"] = tool_call.function.name
-                            tool_call=delta,
+                            delta = f"{buffer['name']}("
-                            parse_status=ToolCallParseStatus.in_progress,
+                            buffer["content"] += delta
-                        ),
+
-                        logprobs=_convert_openai_logprobs(logprobs),
+                        if tool_call.function.arguments:
-                    )
+                            delta = tool_call.function.arguments
-                )
+                            buffer["arguments"] += delta
-        else:
+                            buffer["content"] += delta
+
+                        yield ChatCompletionResponseStreamChunk(
+                            event=ChatCompletionResponseEvent(
+                                event_type=event_type,
+                                delta=ToolCallDelta(
+                                    tool_call=delta,
+                                    parse_status=ToolCallParseStatus.in_progress,
+                                ),
+                                logprobs=_convert_openai_logprobs(logprobs),
+                            )
+                        )
+        elif choice.delta.content:
             yield ChatCompletionResponseStreamChunk(
                 event=ChatCompletionResponseEvent(
                     event_type=event_type,
@@ -935,47 +949,51 @@ async def convert_openai_chat_completion_stream(
                 )
             )
 
-    if toolcall_buffer:
+    for idx, buffer in tool_call_idx_to_buffer.items():
-        delta = ")"
+        logger.debug(f"toolcall_buffer[{idx}]: {buffer}")
-        toolcall_buffer["content"] += delta
+        if buffer["name"]:
-        yield ChatCompletionResponseStreamChunk(
+            delta = ")"
-            event=ChatCompletionResponseEvent(
+            buffer["content"] += delta
-                event_type=event_type,
-                delta=ToolCallDelta(
-                    tool_call=delta,
-                    parse_status=ToolCallParseStatus.in_progress,
-                ),
-                logprobs=_convert_openai_logprobs(logprobs),
-            )
-        )
-        try:
-            arguments = json.loads(toolcall_buffer["arguments"])
-            tool_call = ToolCall(
-                call_id=toolcall_buffer["call_id"],
-                tool_name=toolcall_buffer["name"],
-                arguments=arguments,
-            )
             yield ChatCompletionResponseStreamChunk(
                 event=ChatCompletionResponseEvent(
-                    event_type=ChatCompletionResponseEventType.progress,
+                    event_type=event_type,
                     delta=ToolCallDelta(
-                        tool_call=tool_call,
+                        tool_call=delta,
-                        parse_status=ToolCallParseStatus.succeeded,
+                        parse_status=ToolCallParseStatus.in_progress,
                     ),
-                    stop_reason=stop_reason,
+                    logprobs=None,
                 )
             )
-        except json.JSONDecodeError:
+
-            yield ChatCompletionResponseStreamChunk(
+            try:
-                event=ChatCompletionResponseEvent(
+                arguments = json.loads(buffer["arguments"])
-                    event_type=ChatCompletionResponseEventType.complete,
+                tool_call = ToolCall(
-                    delta=ToolCallDelta(
+                    call_id=buffer["call_id"],
-                        tool_call=toolcall_buffer["content"],
+                    tool_name=buffer["name"],
-                        parse_status=ToolCallParseStatus.failed,
+                    arguments=arguments,
-                    ),
+                )
-                    stop_reason=stop_reason,
+                yield ChatCompletionResponseStreamChunk(
+                    event=ChatCompletionResponseEvent(
+                        event_type=ChatCompletionResponseEventType.progress,
+                        delta=ToolCallDelta(
+                            tool_call=tool_call,
+                            parse_status=ToolCallParseStatus.succeeded,
+                        ),
+                        stop_reason=stop_reason,
+                    )
+                )
+            except json.JSONDecodeError as e:
+                print(f"Failed to parse arguments: {e}")
+                yield ChatCompletionResponseStreamChunk(
+                    event=ChatCompletionResponseEvent(
+                        event_type=ChatCompletionResponseEventType.progress,
+                        delta=ToolCallDelta(
+                            tool_call=buffer["content"],
+                            parse_status=ToolCallParseStatus.failed,
+                        ),
+                        stop_reason=stop_reason,
+                    )
                 )
-            )
 
     yield ChatCompletionResponseStreamChunk(
         event=ChatCompletionResponseEvent(

llama_stack/scripts/run_tests.sh

@@ -1,15 +0,0 @@
-#!/bin/bash
-
-# 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.
-
-THIS_DIR="$(cd "$(dirname "$(readlink -f "${BASH_SOURCE[0]}")")" && pwd)"
-
-set -euo pipefail
-set -x
-
-stack_dir=$(dirname $(dirname $THIS_DIR))
-PYTHONPATH=$stack_dir pytest -p no:warnings --asyncio-mode auto --tb=short

llama_stack/templates/nvidia/build.yaml

@@ -1,13 +1,13 @@
 version: '2'
 distribution_spec:
-  description: Use NVIDIA NIM for running LLM inference
+  description: Use NVIDIA NIM for running LLM inference and safety
   providers:
     inference:
     - remote::nvidia
     vector_io:
     - inline::faiss
     safety:
-    - inline::llama-guard
+    - remote::nvidia
     agents:
     - inline::meta-reference
     telemetry:
@@ -15,16 +15,9 @@ distribution_spec:
     eval:
     - inline::meta-reference
     datasetio:
-    - remote::huggingface
     - inline::localfs
     scoring:
     - inline::basic
-    - inline::llm-as-judge
-    - inline::braintrust
     tool_runtime:
-    - remote::brave-search
-    - remote::tavily-search
-    - inline::code-interpreter
     - inline::rag-runtime
-    - remote::model-context-protocol
 image_type: conda

llama_stack/templates/nvidia/nvidia.py

@@ -6,9 +6,10 @@
 
 from pathlib import Path
 
-from llama_stack.distribution.datatypes import Provider, ToolGroupInput
+from llama_stack.distribution.datatypes import ModelInput, Provider, ShieldInput, ToolGroupInput
 from llama_stack.providers.remote.inference.nvidia import NVIDIAConfig
 from llama_stack.providers.remote.inference.nvidia.models import MODEL_ENTRIES
+from llama_stack.providers.remote.safety.nvidia import NVIDIASafetyConfig
 from llama_stack.templates.template import DistributionTemplate, RunConfigSettings, get_model_registry
 
 
@@ -16,19 +17,13 @@ def get_distribution_template() -> DistributionTemplate:
     providers = {
         "inference": ["remote::nvidia"],
         "vector_io": ["inline::faiss"],
-        "safety": ["inline::llama-guard"],
+        "safety": ["remote::nvidia"],
         "agents": ["inline::meta-reference"],
         "telemetry": ["inline::meta-reference"],
         "eval": ["inline::meta-reference"],
-        "datasetio": ["remote::huggingface", "inline::localfs"],
+        "datasetio": ["inline::localfs"],
-        "scoring": ["inline::basic", "inline::llm-as-judge", "inline::braintrust"],
+        "scoring": ["inline::basic"],
-        "tool_runtime": [
+        "tool_runtime": ["inline::rag-runtime"],
-            "remote::brave-search",
-            "remote::tavily-search",
-            "inline::code-interpreter",
-            "inline::rag-runtime",
-            "remote::model-context-protocol",
-        ],
     }
 
     inference_provider = Provider(
@@ -36,30 +31,35 @@ def get_distribution_template() -> DistributionTemplate:
         provider_type="remote::nvidia",
         config=NVIDIAConfig.sample_run_config(),
     )
+    safety_provider = Provider(
+        provider_id="nvidia",
+        provider_type="remote::nvidia",
+        config=NVIDIASafetyConfig.sample_run_config(),
+    )
+    inference_model = ModelInput(
+        model_id="${env.INFERENCE_MODEL}",
+        provider_id="nvidia",
+    )
+    safety_model = ModelInput(
+        model_id="${env.SAFETY_MODEL}",
+        provider_id="nvidia",
+    )
 
     available_models = {
         "nvidia": MODEL_ENTRIES,
     }
     default_tool_groups = [
-        ToolGroupInput(
-            toolgroup_id="builtin::websearch",
-            provider_id="tavily-search",
-        ),
         ToolGroupInput(
             toolgroup_id="builtin::rag",
             provider_id="rag-runtime",
         ),
-        ToolGroupInput(
-            toolgroup_id="builtin::code_interpreter",
-            provider_id="code-interpreter",
-        ),
     ]
 
     default_models = get_model_registry(available_models)
     return DistributionTemplate(
         name="nvidia",
         distro_type="remote_hosted",
-        description="Use NVIDIA NIM for running LLM inference",
+        description="Use NVIDIA NIM for running LLM inference and safety",
         container_image=None,
         template_path=Path(__file__).parent / "doc_template.md",
         providers=providers,
@@ -72,15 +72,34 @@ def get_distribution_template() -> DistributionTemplate:
                 default_models=default_models,
                 default_tool_groups=default_tool_groups,
             ),
+            "run-with-safety.yaml": RunConfigSettings(
+                provider_overrides={
+                    "inference": [
+                        inference_provider,
+                        safety_provider,
+                    ]
+                },
+                default_models=[inference_model, safety_model],
+                default_shields=[ShieldInput(shield_id="${env.SAFETY_MODEL}", provider_id="nvidia")],
+                default_tool_groups=default_tool_groups,
+            ),
         },
         run_config_env_vars={
-            "LLAMASTACK_PORT": (
-                "5001",
-                "Port for the Llama Stack distribution server",
-            ),
             "NVIDIA_API_KEY": (
                 "",
                 "NVIDIA API Key",
             ),
+            "GUARDRAILS_SERVICE_URL": (
+                "http://0.0.0.0:7331",
+                "URL for the NeMo Guardrails Service",
+            ),
+            "INFERENCE_MODEL": (
+                "Llama3.1-8B-Instruct",
+                "Inference model",
+            ),
+            "SAFETY_MODEL": (
+                "meta/llama-3.1-8b-instruct",
+                "Name of the model to use for safety",
+            ),
         },
     )

llama_stack/templates/nvidia/run-with-safety.yaml

@@ -0,0 +1,101 @@
+version: '2'
+image_name: nvidia
+apis:
+- agents
+- datasetio
+- eval
+- inference
+- safety
+- scoring
+- telemetry
+- tool_runtime
+- vector_io
+providers:
+  inference:
+  - provider_id: nvidia
+    provider_type: remote::nvidia
+    config:
+      url: ${env.NVIDIA_BASE_URL:https://integrate.api.nvidia.com}
+      api_key: ${env.NVIDIA_API_KEY:}
+  - provider_id: nvidia
+    provider_type: remote::nvidia
+    config:
+      guardrails_service_url: ${env.GUARDRAILS_SERVICE_URL:http://localhost:7331}
+      config_id: self-check
+  vector_io:
+  - provider_id: faiss
+    provider_type: inline::faiss
+    config:
+      kvstore:
+        type: sqlite
+        namespace: null
+        db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/nvidia}/faiss_store.db
+  safety:
+  - provider_id: nvidia
+    provider_type: remote::nvidia
+    config:
+      guardrails_service_url: ${env.GUARDRAILS_SERVICE_URL:http://localhost:7331}
+      config_id: self-check
+  agents:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config:
+      persistence_store:
+        type: sqlite
+        namespace: null
+        db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/nvidia}/agents_store.db
+  telemetry:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config:
+      service_name: ${env.OTEL_SERVICE_NAME:llama-stack}
+      sinks: ${env.TELEMETRY_SINKS:console,sqlite}
+      sqlite_db_path: ${env.SQLITE_DB_PATH:~/.llama/distributions/nvidia/trace_store.db}
+  eval:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config:
+      kvstore:
+        type: sqlite
+        namespace: null
+        db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/nvidia}/meta_reference_eval.db
+  datasetio:
+  - provider_id: localfs
+    provider_type: inline::localfs
+    config:
+      kvstore:
+        type: sqlite
+        namespace: null
+        db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/nvidia}/localfs_datasetio.db
+  scoring:
+  - provider_id: basic
+    provider_type: inline::basic
+    config: {}
+  tool_runtime:
+  - provider_id: rag-runtime
+    provider_type: inline::rag-runtime
+    config: {}
+metadata_store:
+  type: sqlite
+  db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/nvidia}/registry.db
+models:
+- metadata: {}
+  model_id: ${env.INFERENCE_MODEL}
+  provider_id: nvidia
+  model_type: llm
+- metadata: {}
+  model_id: ${env.SAFETY_MODEL}
+  provider_id: nvidia
+  model_type: llm
+shields:
+- shield_id: ${env.SAFETY_MODEL}
+  provider_id: nvidia
+vector_dbs: []
+datasets: []
+scoring_fns: []
+benchmarks: []
+tool_groups:
+- toolgroup_id: builtin::rag
+  provider_id: rag-runtime
+server:
+  port: 8321

llama_stack/templates/nvidia/run.yaml

@@ -26,10 +26,11 @@ providers:
         namespace: null
         db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/nvidia}/faiss_store.db
   safety:
-  - provider_id: llama-guard
+  - provider_id: nvidia
-    provider_type: inline::llama-guard
+    provider_type: remote::nvidia
     config:
-      excluded_categories: []
+      guardrails_service_url: ${env.GUARDRAILS_SERVICE_URL:http://localhost:7331}
+      config_id: self-check
   agents:
   - provider_id: meta-reference
     provider_type: inline::meta-reference
@@ -54,13 +55,6 @@ providers:
         namespace: null
         db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/nvidia}/meta_reference_eval.db
   datasetio:
-  - provider_id: huggingface
-    provider_type: remote::huggingface
-    config:
-      kvstore:
-        type: sqlite
-        namespace: null
-        db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/nvidia}/huggingface_datasetio.db
   - provider_id: localfs
     provider_type: inline::localfs
     config:
@@ -72,33 +66,10 @@ providers:
   - provider_id: basic
     provider_type: inline::basic
     config: {}
-  - provider_id: llm-as-judge
-    provider_type: inline::llm-as-judge
-    config: {}
-  - provider_id: braintrust
-    provider_type: inline::braintrust
-    config:
-      openai_api_key: ${env.OPENAI_API_KEY:}
   tool_runtime:
-  - provider_id: brave-search
-    provider_type: remote::brave-search
-    config:
-      api_key: ${env.BRAVE_SEARCH_API_KEY:}
-      max_results: 3
-  - provider_id: tavily-search
-    provider_type: remote::tavily-search
-    config:
-      api_key: ${env.TAVILY_SEARCH_API_KEY:}
-      max_results: 3
-  - provider_id: code-interpreter
-    provider_type: inline::code-interpreter
-    config: {}
   - provider_id: rag-runtime
     provider_type: inline::rag-runtime
     config: {}
-  - provider_id: model-context-protocol
-    provider_type: remote::model-context-protocol
-    config: {}
 metadata_store:
   type: sqlite
   db_path: ${env.SQLITE_STORE_DIR:~/.llama/distributions/nvidia}/registry.db
@@ -227,11 +198,7 @@ datasets: []
 scoring_fns: []
 benchmarks: []
 tool_groups:
-- toolgroup_id: builtin::websearch
-  provider_id: tavily-search
 - toolgroup_id: builtin::rag
   provider_id: rag-runtime
-- toolgroup_id: builtin::code_interpreter
-  provider_id: code-interpreter
 server:
   port: 8321

llama_stack/templates/open-benchmark/open_benchmark.py

@@ -6,7 +6,7 @@
 
 from typing import Dict, List, Tuple
 
-from llama_stack.apis.common.content_types import URL
+from llama_stack.apis.datasets import DatasetPurpose, URIDataSource
 from llama_stack.apis.models.models import ModelType
 from llama_stack.distribution.datatypes import (
     BenchmarkInput,
@@ -171,76 +171,42 @@ def get_distribution_template() -> DistributionTemplate:
         DatasetInput(
             dataset_id="simpleqa",
             provider_id="huggingface",
-            url=URL(uri="https://huggingface.co/datasets/llamastack/simpleqa"),
+            purpose=DatasetPurpose.eval_messages_answer,
-            metadata={
+            source=URIDataSource(
-                "path": "llamastack/simpleqa",
+                uri="huggingface://datasets/llamastack/simpleqa?split=train",
-                "split": "train",
+            ),
-            },
-            dataset_schema={
-                "input_query": {"type": "string"},
-                "expected_answer": {"type": "string"},
-                "chat_completion_input": {"type": "string"},
-            },
         ),
         DatasetInput(
             dataset_id="mmlu_cot",
             provider_id="huggingface",
-            url=URL(uri="https://huggingface.co/datasets/llamastack/mmlu_cot"),
+            purpose=DatasetPurpose.eval_messages_answer,
-            metadata={
+            source=URIDataSource(
-                "path": "llamastack/mmlu_cot",
+                uri="huggingface://datasets/llamastack/mmlu_cot?split=test&name=all",
-                "name": "all",
+            ),
-                "split": "test",
-            },
-            dataset_schema={
-                "input_query": {"type": "string"},
-                "expected_answer": {"type": "string"},
-                "chat_completion_input": {"type": "string"},
-            },
         ),
         DatasetInput(
             dataset_id="gpqa_cot",
             provider_id="huggingface",
-            url=URL(uri="https://huggingface.co/datasets/llamastack/gpqa_0shot_cot"),
+            purpose=DatasetPurpose.eval_messages_answer,
-            metadata={
+            source=URIDataSource(
-                "path": "llamastack/gpqa_0shot_cot",
+                uri="huggingface://datasets/llamastack/gpqa_0shot_cot?split=test&name=gpqa_main",
-                "name": "gpqa_main",
+            ),
-                "split": "train",
-            },
-            dataset_schema={
-                "input_query": {"type": "string"},
-                "expected_answer": {"type": "string"},
-                "chat_completion_input": {"type": "string"},
-            },
         ),
         DatasetInput(
             dataset_id="math_500",
             provider_id="huggingface",
-            url=URL(uri="https://huggingface.co/datasets/llamastack/math_500"),
+            purpose=DatasetPurpose.eval_messages_answer,
-            metadata={
+            source=URIDataSource(
-                "path": "llamastack/math_500",
+                uri="huggingface://datasets/llamastack/math_500?split=test",
-                "split": "test",
+            ),
-            },
-            dataset_schema={
-                "input_query": {"type": "string"},
-                "expected_answer": {"type": "string"},
-                "chat_completion_input": {"type": "string"},
-            },
         ),
         DatasetInput(
             dataset_id="bfcl",
             provider_id="huggingface",
-            url=URL(uri="https://huggingface.co/datasets/llamastack/bfcl_v3"),
+            purpose=DatasetPurpose.eval_messages_answer,
-            metadata={
+            source=URIDataSource(
-                "path": "llamastack/bfcl_v3",
+                uri="huggingface://datasets/llamastack/bfcl_v3?split=train",
-                "split": "train",
+            ),
-            },
-            dataset_schema={
-                "function": {"type": "string"},
-                "language": {"type": "string"},
-                "ground_truth": {"type": "string"},
-                "id": {"type": "string"},
-                "chat_completion_input": {"type": "string"},
-            },
         ),
     ]
 

llama_stack/templates/open-benchmark/run.yaml

@@ -158,80 +158,39 @@ shields:
 - shield_id: meta-llama/Llama-Guard-3-8B
 vector_dbs: []
 datasets:
-- dataset_schema:
+- purpose: eval/messages-answer
-    input_query:
+  source:
-      type: string
+    type: uri
-    expected_answer:
+    uri: huggingface://datasets/llamastack/simpleqa?split=train
-      type: string
+  metadata: {}
-    chat_completion_input:
-      type: string
-  url:
-    uri: https://huggingface.co/datasets/llamastack/simpleqa
-  metadata:
-    path: llamastack/simpleqa
-    split: train
   dataset_id: simpleqa
   provider_id: huggingface
-- dataset_schema:
+- purpose: eval/messages-answer
-    input_query:
+  source:
-      type: string
+    type: uri
-    expected_answer:
+    uri: huggingface://datasets/llamastack/mmlu_cot?split=test&name=all
-      type: string
+  metadata: {}
-    chat_completion_input:
-      type: string
-  url:
-    uri: https://huggingface.co/datasets/llamastack/mmlu_cot
-  metadata:
-    path: llamastack/mmlu_cot
-    name: all
-    split: test
   dataset_id: mmlu_cot
   provider_id: huggingface
-- dataset_schema:
+- purpose: eval/messages-answer
-    input_query:
+  source:
-      type: string
+    type: uri
-    expected_answer:
+    uri: huggingface://datasets/llamastack/gpqa_0shot_cot?split=test&name=gpqa_main
-      type: string
+  metadata: {}
-    chat_completion_input:
-      type: string
-  url:
-    uri: https://huggingface.co/datasets/llamastack/gpqa_0shot_cot
-  metadata:
-    path: llamastack/gpqa_0shot_cot
-    name: gpqa_main
-    split: train
   dataset_id: gpqa_cot
   provider_id: huggingface
-- dataset_schema:
+- purpose: eval/messages-answer
-    input_query:
+  source:
-      type: string
+    type: uri
-    expected_answer:
+    uri: huggingface://datasets/llamastack/math_500?split=test
-      type: string
+  metadata: {}
-    chat_completion_input:
-      type: string
-  url:
-    uri: https://huggingface.co/datasets/llamastack/math_500
-  metadata:
-    path: llamastack/math_500
-    split: test
   dataset_id: math_500
   provider_id: huggingface
-- dataset_schema:
+- purpose: eval/messages-answer
-    function:
+  source:
-      type: string
+    type: uri
-    language:
+    uri: huggingface://datasets/llamastack/bfcl_v3?split=train
-      type: string
+  metadata: {}
-    ground_truth:
-      type: string
-    id:
-      type: string
-    chat_completion_input:
-      type: string
-  url:
-    uri: https://huggingface.co/datasets/llamastack/bfcl_v3
-  metadata:
-    path: llamastack/bfcl_v3
-    split: train
   dataset_id: bfcl
   provider_id: huggingface
 scoring_fns: []

llama_stack/templates/remote-vllm/run-with-safety.yaml

@@ -18,12 +18,14 @@ providers:
       url: ${env.VLLM_URL}
       max_tokens: ${env.VLLM_MAX_TOKENS:4096}
       api_token: ${env.VLLM_API_TOKEN:fake}
+      tls_verify: ${env.VLLM_TLS_VERIFY:true}
   - provider_id: vllm-safety
     provider_type: remote::vllm
     config:
       url: ${env.SAFETY_VLLM_URL}
       max_tokens: ${env.VLLM_MAX_TOKENS:4096}
       api_token: ${env.VLLM_API_TOKEN:fake}
+      tls_verify: ${env.VLLM_TLS_VERIFY:true}
   - provider_id: sentence-transformers
     provider_type: inline::sentence-transformers
     config: {}

llama_stack/templates/remote-vllm/run.yaml

@@ -18,6 +18,7 @@ providers:
       url: ${env.VLLM_URL}
       max_tokens: ${env.VLLM_MAX_TOKENS:4096}
       api_token: ${env.VLLM_API_TOKEN:fake}
+      tls_verify: ${env.VLLM_TLS_VERIFY:true}
   - provider_id: sentence-transformers
     provider_type: inline::sentence-transformers
     config: {}

llama_stack/templates/template.py

@@ -11,6 +11,7 @@ import jinja2
 import yaml
 from pydantic import BaseModel, Field
 
+from llama_stack.apis.datasets import DatasetPurpose
 from llama_stack.apis.models.models import ModelType
 from llama_stack.distribution.datatypes import (
     Api,
@@ -214,7 +215,9 @@ class DistributionTemplate(BaseModel):
 
         # Register YAML representer for ModelType
         yaml.add_representer(ModelType, enum_representer)
+        yaml.add_representer(DatasetPurpose, enum_representer)
         yaml.SafeDumper.add_representer(ModelType, enum_representer)
+        yaml.SafeDumper.add_representer(DatasetPurpose, enum_representer)
 
         for output_dir in [yaml_output_dir, doc_output_dir]:
             output_dir.mkdir(parents=True, exist_ok=True)

pyproject.toml

@@ -269,6 +269,7 @@ exclude = [
     "^llama_stack/providers/remote/inference/together/",
     "^llama_stack/providers/remote/inference/vllm/",
     "^llama_stack/providers/remote/safety/bedrock/",
+    "^llama_stack/providers/remote/safety/nvidia/",
     "^llama_stack/providers/remote/safety/sample/",
     "^llama_stack/providers/remote/tool_runtime/bing_search/",
     "^llama_stack/providers/remote/tool_runtime/brave_search/",

scripts/distro_codegen.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -20,7 +21,7 @@ from llama_stack.distribution.build import (
     get_provider_dependencies,
 )
 
-REPO_ROOT = Path(__file__).parent.parent.parent
+REPO_ROOT = Path(__file__).parent.parent
 
 
 class ChangedPathTracker:

scripts/gen-changelog.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #

scripts/generate_prompt_format.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -18,7 +19,7 @@ import fire
 
 from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.providers.inline.inference.meta_reference.config import MetaReferenceInferenceConfig
-from llama_stack.providers.inline.inference.meta_reference.generation import Llama
+from llama_stack.providers.inline.inference.meta_reference.llama3.generation import Llama3
 
 THIS_DIR = Path(__file__).parent.resolve()
 
@@ -41,7 +42,7 @@ def run_main(
     llama_model = resolve_model(model_id)
     if not llama_model:
         raise ValueError(f"Model {model_id} not found")
-    generator = Llama.build(
+    generator = Llama3.build(
         config=config,
         model_id=model_id,
         llama_model=llama_model,

scripts/run_client_sdk_tests.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -15,8 +16,7 @@ Script for running api on AsyncLlamaStackAsLibraryClient with templates
 
 Assuming directory structure:
 - llama-stack
-    - llama_stack
+    - scripts
-        - scripts
     - tests
         - api
 
@@ -25,10 +25,10 @@ Example command:
 cd llama-stack
 EXPORT TOGETHER_API_KEY=<..>
 EXPORT FIREWORKS_API_KEY=<..>
-python llama_stack/scripts/run_client_sdk_tests.py --templates together fireworks --report
+./scripts/run_client_sdk_tests.py --templates together fireworks --report
 """
 
-REPO_ROOT = Path(__file__).parent.parent.parent
+REPO_ROOT = Path(__file__).parent.parent
 CLIENT_SDK_TESTS_RELATIVE_PATH = "tests/api/"
 
 

tests/integration/agents/test_agents.py

@@ -271,7 +271,7 @@ def test_custom_tool(llama_stack_client_with_mocked_inference, agent_config):
     client_tool = get_boiling_point
     agent_config = {
         **agent_config,
-        "tools": ["builtin::websearch", client_tool],
+        "tools": [client_tool],
     }
 
     agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
@@ -320,42 +320,55 @@ def test_custom_tool_infinite_loop(llama_stack_client_with_mocked_inference, age
     assert num_tool_calls <= 5
 
 
-def test_tool_choice(llama_stack_client_with_mocked_inference, agent_config):
+def test_tool_choice_required(llama_stack_client_with_mocked_inference, agent_config):
-    def run_agent(tool_choice):
+    tool_execution_steps = run_agent_with_tool_choice(
-        client_tool = get_boiling_point
+        llama_stack_client_with_mocked_inference, agent_config, "required"
-
+    )
-        test_agent_config = {
-            **agent_config,
-            "tool_config": {"tool_choice": tool_choice},
-            "tools": [client_tool],
-        }
-
-        agent = Agent(llama_stack_client_with_mocked_inference, **test_agent_config)
-        session_id = agent.create_session(f"test-session-{uuid4()}")
-
-        response = agent.create_turn(
-            messages=[
-                {
-                    "role": "user",
-                    "content": "What is the boiling point of polyjuice?",
-                },
-            ],
-            session_id=session_id,
-            stream=False,
-        )
-
-        return [step for step in response.steps if step.step_type == "tool_execution"]
-
-    tool_execution_steps = run_agent("required")
     assert len(tool_execution_steps) > 0
 
-    tool_execution_steps = run_agent("none")
+
+def test_tool_choice_none(llama_stack_client_with_mocked_inference, agent_config):
+    tool_execution_steps = run_agent_with_tool_choice(llama_stack_client_with_mocked_inference, agent_config, "none")
     assert len(tool_execution_steps) == 0
 
-    tool_execution_steps = run_agent("get_boiling_point")
+
+def test_tool_choice_get_boiling_point(llama_stack_client_with_mocked_inference, agent_config):
+    if "llama" not in agent_config["model"].lower():
+        pytest.xfail("NotImplemented for non-llama models")
+
+    tool_execution_steps = run_agent_with_tool_choice(
+        llama_stack_client_with_mocked_inference, agent_config, "get_boiling_point"
+    )
     assert len(tool_execution_steps) >= 1 and tool_execution_steps[0].tool_calls[0].tool_name == "get_boiling_point"
 
 
+def run_agent_with_tool_choice(client, agent_config, tool_choice):
+    client_tool = get_boiling_point
+
+    test_agent_config = {
+        **agent_config,
+        "tool_config": {"tool_choice": tool_choice},
+        "tools": [client_tool],
+        "max_infer_iters": 2,
+    }
+
+    agent = Agent(client, **test_agent_config)
+    session_id = agent.create_session(f"test-session-{uuid4()}")
+
+    response = agent.create_turn(
+        messages=[
+            {
+                "role": "user",
+                "content": "What is the boiling point of polyjuice?",
+            },
+        ],
+        session_id=session_id,
+        stream=False,
+    )
+
+    return [step for step in response.steps if step.step_type == "tool_execution"]
+
+
 @pytest.mark.parametrize("rag_tool_name", ["builtin::rag/knowledge_search", "builtin::rag"])
 def test_rag_agent(llama_stack_client_with_mocked_inference, agent_config, rag_tool_name):
     urls = ["chat.rst", "llama3.rst", "memory_optimizations.rst", "lora_finetune.rst"]
@@ -571,7 +584,7 @@ def test_rag_and_code_agent(llama_stack_client_with_mocked_inference, agent_conf
     [(get_boiling_point, False), (get_boiling_point_with_metadata, True)],
 )
 def test_create_turn_response(llama_stack_client_with_mocked_inference, agent_config, client_tools):
-    client_tool, expectes_metadata = client_tools
+    client_tool, expects_metadata = client_tools
     agent_config = {
         **agent_config,
         "input_shields": [],
@@ -597,7 +610,7 @@ def test_create_turn_response(llama_stack_client_with_mocked_inference, agent_co
     assert steps[0].step_type == "inference"
     assert steps[1].step_type == "tool_execution"
     assert steps[1].tool_calls[0].tool_name.startswith("get_boiling_point")
-    if expectes_metadata:
+    if expects_metadata:
         assert steps[1].tool_responses[0].metadata["source"] == "https://www.google.com"
     assert steps[2].step_type == "inference"
 
@@ -609,3 +622,44 @@ def test_create_turn_response(llama_stack_client_with_mocked_inference, agent_co
             assert last_step_completed_at < step.started_at
             assert step.started_at < step.completed_at
             last_step_completed_at = step.completed_at
+
+
+def test_multi_tool_calls(llama_stack_client_with_mocked_inference, agent_config):
+    if "gpt" not in agent_config["model"]:
+        pytest.xfail("Only tested on GPT models")
+
+    agent_config = {
+        **agent_config,
+        "tools": [get_boiling_point],
+    }
+
+    agent = Agent(llama_stack_client_with_mocked_inference, **agent_config)
+    session_id = agent.create_session(f"test-session-{uuid4()}")
+
+    response = agent.create_turn(
+        messages=[
+            {
+                "role": "user",
+                "content": "Call get_boiling_point twice to answer: What is the boiling point of polyjuice in both celsius and fahrenheit?",
+            },
+        ],
+        session_id=session_id,
+        stream=False,
+    )
+    steps = response.steps
+    assert len(steps) == 7
+    assert steps[0].step_type == "shield_call"
+    assert steps[1].step_type == "inference"
+    assert steps[2].step_type == "shield_call"
+    assert steps[3].step_type == "tool_execution"
+    assert steps[4].step_type == "shield_call"
+    assert steps[5].step_type == "inference"
+    assert steps[6].step_type == "shield_call"
+
+    tool_execution_step = steps[3]
+    assert len(tool_execution_step.tool_calls) == 2
+    assert tool_execution_step.tool_calls[0].tool_name.startswith("get_boiling_point")
+    assert tool_execution_step.tool_calls[1].tool_name.startswith("get_boiling_point")
+
+    output = response.output_message.content.lower()
+    assert "-100" in output and "-212" in output

tests/integration/datasetio/__init__.py

@@ -1,5 +0,0 @@
-# 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.

tests/integration/datasetio/test_datasetio.py

@@ -1,114 +0,0 @@
-# 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 base64
-import mimetypes
-import os
-from pathlib import Path
-
-import pytest
-
-# How to run this test:
-#
-# LLAMA_STACK_CONFIG="template-name" pytest -v tests/integration/datasetio
-
-
-@pytest.fixture
-def dataset_for_test(llama_stack_client):
-    dataset_id = "test_dataset"
-    register_dataset(llama_stack_client, dataset_id=dataset_id)
-    yield
-    # Teardown - this always runs, even if the test fails
-    try:
-        llama_stack_client.datasets.unregister(dataset_id)
-    except Exception as e:
-        print(f"Warning: Failed to unregister test_dataset: {e}")
-
-
-def data_url_from_file(file_path: str) -> str:
-    if not os.path.exists(file_path):
-        raise FileNotFoundError(f"File not found: {file_path}")
-
-    with open(file_path, "rb") as file:
-        file_content = file.read()
-
-    base64_content = base64.b64encode(file_content).decode("utf-8")
-    mime_type, _ = mimetypes.guess_type(file_path)
-
-    data_url = f"data:{mime_type};base64,{base64_content}"
-
-    return data_url
-
-
-def register_dataset(llama_stack_client, for_generation=False, for_rag=False, dataset_id="test_dataset"):
-    if for_rag:
-        test_file = Path(os.path.abspath(__file__)).parent / "test_rag_dataset.csv"
-    else:
-        test_file = Path(os.path.abspath(__file__)).parent / "test_dataset.csv"
-    test_url = data_url_from_file(str(test_file))
-
-    if for_generation:
-        dataset_schema = {
-            "expected_answer": {"type": "string"},
-            "input_query": {"type": "string"},
-            "chat_completion_input": {"type": "chat_completion_input"},
-        }
-    elif for_rag:
-        dataset_schema = {
-            "expected_answer": {"type": "string"},
-            "input_query": {"type": "string"},
-            "generated_answer": {"type": "string"},
-            "context": {"type": "string"},
-        }
-    else:
-        dataset_schema = {
-            "expected_answer": {"type": "string"},
-            "input_query": {"type": "string"},
-            "generated_answer": {"type": "string"},
-        }
-
-    dataset_providers = [x for x in llama_stack_client.providers.list() if x.api == "datasetio"]
-    dataset_provider_id = dataset_providers[0].provider_id
-
-    llama_stack_client.datasets.register(
-        dataset_id=dataset_id,
-        dataset_schema=dataset_schema,
-        url=dict(uri=test_url),
-        provider_id=dataset_provider_id,
-    )
-
-
-def test_register_unregister_dataset(llama_stack_client):
-    register_dataset(llama_stack_client)
-    response = llama_stack_client.datasets.list()
-    assert isinstance(response, list)
-    assert len(response) == 1
-    assert response[0].identifier == "test_dataset"
-
-    llama_stack_client.datasets.unregister("test_dataset")
-    response = llama_stack_client.datasets.list()
-    assert isinstance(response, list)
-    assert len(response) == 0
-
-
-def test_get_rows_paginated(llama_stack_client, dataset_for_test):
-    response = llama_stack_client.datasetio.get_rows_paginated(
-        dataset_id="test_dataset",
-        rows_in_page=3,
-    )
-    assert isinstance(response.rows, list)
-    assert len(response.rows) == 3
-    assert response.next_page_token == "3"
-
-    # iterate over all rows
-    response = llama_stack_client.datasetio.get_rows_paginated(
-        dataset_id="test_dataset",
-        rows_in_page=2,
-        page_token=response.next_page_token,
-    )
-    assert isinstance(response.rows, list)
-    assert len(response.rows) == 2
-    assert response.next_page_token == "5"

tests/integration/datasets/test_datasets.py

@@ -0,0 +1,95 @@
+# 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 base64
+import mimetypes
+import os
+
+import pytest
+
+# How to run this test:
+#
+# LLAMA_STACK_CONFIG="template-name" pytest -v tests/integration/datasets
+
+
+def data_url_from_file(file_path: str) -> str:
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    with open(file_path, "rb") as file:
+        file_content = file.read()
+
+    base64_content = base64.b64encode(file_content).decode("utf-8")
+    mime_type, _ = mimetypes.guess_type(file_path)
+
+    data_url = f"data:{mime_type};base64,{base64_content}"
+
+    return data_url
+
+
+@pytest.mark.parametrize(
+    "purpose, source, provider_id, limit",
+    [
+        (
+            "eval/messages-answer",
+            {
+                "type": "uri",
+                "uri": "huggingface://datasets/llamastack/simpleqa?split=train",
+            },
+            "huggingface",
+            10,
+        ),
+        (
+            "eval/messages-answer",
+            {
+                "type": "rows",
+                "rows": [
+                    {
+                        "messages": [{"role": "user", "content": "Hello, world!"}],
+                        "answer": "Hello, world!",
+                    },
+                    {
+                        "messages": [
+                            {
+                                "role": "user",
+                                "content": "What is the capital of France?",
+                            }
+                        ],
+                        "answer": "Paris",
+                    },
+                ],
+            },
+            "localfs",
+            2,
+        ),
+        (
+            "eval/messages-answer",
+            {
+                "type": "uri",
+                "uri": data_url_from_file(os.path.join(os.path.dirname(__file__), "test_dataset.csv")),
+            },
+            "localfs",
+            5,
+        ),
+    ],
+)
+def test_register_and_iterrows(llama_stack_client, purpose, source, provider_id, limit):
+    dataset = llama_stack_client.datasets.register(
+        purpose=purpose,
+        source=source,
+    )
+    assert dataset.identifier is not None
+    assert dataset.provider_id == provider_id
+    iterrow_response = llama_stack_client.datasets.iterrows(dataset.identifier, limit=limit)
+    assert len(iterrow_response.data) == limit
+
+    dataset_list = llama_stack_client.datasets.list()
+    assert dataset.identifier in [d.identifier for d in dataset_list]
+
+    llama_stack_client.datasets.unregister(dataset.identifier)
+    dataset_list = llama_stack_client.datasets.list()
+    assert dataset.identifier not in [d.identifier for d in dataset_list]

tests/integration/eval/test_eval.py

@@ -4,10 +4,11 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 import uuid
+from pathlib import Path
 
 import pytest
 
-from ..datasetio.test_datasetio import register_dataset
+from ..datasets.test_datasets import data_url_from_file
 
 # How to run this test:
 #
@@ -16,15 +17,21 @@ from ..datasetio.test_datasetio import register_dataset
 
 @pytest.mark.parametrize("scoring_fn_id", ["basic::equality"])
 def test_evaluate_rows(llama_stack_client, text_model_id, scoring_fn_id):
-    register_dataset(llama_stack_client, for_generation=True, dataset_id="test_dataset_for_eval")
+    dataset = llama_stack_client.datasets.register(
-    response = llama_stack_client.datasets.list()
+        purpose="eval/messages-answer",
-    assert any(x.identifier == "test_dataset_for_eval" for x in response)
+        source={
-
+            "type": "uri",
-    rows = llama_stack_client.datasetio.get_rows_paginated(
+            "uri": data_url_from_file(Path(__file__).parent.parent / "datasets" / "test_dataset.csv"),
-        dataset_id="test_dataset_for_eval",
+        },
-        rows_in_page=3,
     )
-    assert len(rows.rows) == 3
+    response = llama_stack_client.datasets.list()
+    assert any(x.identifier == dataset.identifier for x in response)
+
+    rows = llama_stack_client.datasets.iterrows(
+        dataset_id=dataset.identifier,
+        limit=3,
+    )
+    assert len(rows.data) == 3
 
     scoring_functions = [
         scoring_fn_id,
@@ -32,7 +39,7 @@ def test_evaluate_rows(llama_stack_client, text_model_id, scoring_fn_id):
     benchmark_id = str(uuid.uuid4())
     llama_stack_client.benchmarks.register(
         benchmark_id=benchmark_id,
-        dataset_id="test_dataset_for_eval",
+        dataset_id=dataset.identifier,
         scoring_functions=scoring_functions,
     )
     list_benchmarks = llama_stack_client.benchmarks.list()
@@ -40,7 +47,7 @@ def test_evaluate_rows(llama_stack_client, text_model_id, scoring_fn_id):
 
     response = llama_stack_client.eval.evaluate_rows(
         benchmark_id=benchmark_id,
-        input_rows=rows.rows,
+        input_rows=rows.data,
         scoring_functions=scoring_functions,
         benchmark_config={
             "eval_candidate": {
@@ -59,11 +66,17 @@ def test_evaluate_rows(llama_stack_client, text_model_id, scoring_fn_id):
 
 @pytest.mark.parametrize("scoring_fn_id", ["basic::subset_of"])
 def test_evaluate_benchmark(llama_stack_client, text_model_id, scoring_fn_id):
-    register_dataset(llama_stack_client, for_generation=True, dataset_id="test_dataset_for_eval_2")
+    dataset = llama_stack_client.datasets.register(
+        purpose="eval/messages-answer",
+        source={
+            "type": "uri",
+            "uri": data_url_from_file(Path(__file__).parent.parent / "datasets" / "test_dataset.csv"),
+        },
+    )
     benchmark_id = str(uuid.uuid4())
     llama_stack_client.benchmarks.register(
         benchmark_id=benchmark_id,
-        dataset_id="test_dataset_for_eval_2",
+        dataset_id=dataset.identifier,
         scoring_functions=[scoring_fn_id],
     )
 

tests/integration/scoring/test_scoring.py

@@ -5,23 +5,11 @@
 # the root directory of this source tree.
 
 
+from pathlib import Path
+
+import pandas as pd
 import pytest
 
-from ..datasetio.test_datasetio import register_dataset
-
-
-@pytest.fixture
-def rag_dataset_for_test(llama_stack_client):
-    dataset_id = "test_dataset"
-    register_dataset(llama_stack_client, for_rag=True, dataset_id=dataset_id)
-    yield  # This is where the test function will run
-
-    # Teardown - this always runs, even if the test fails
-    try:
-        llama_stack_client.datasets.unregister(dataset_id)
-    except Exception as e:
-        print(f"Warning: Failed to unregister test_dataset: {e}")
-
 
 @pytest.fixture
 def sample_judge_prompt_template():
@@ -92,49 +80,34 @@ def test_scoring_functions_register(
     # TODO: add unregister api for scoring functions
 
 
-def test_scoring_score(llama_stack_client, rag_dataset_for_test):
+@pytest.mark.parametrize("scoring_fn_id", ["basic::equality"])
+def test_scoring_score(llama_stack_client, scoring_fn_id):
     # scoring individual rows
-    rows = llama_stack_client.datasetio.get_rows_paginated(
+    df = pd.read_csv(Path(__file__).parent.parent / "datasets" / "test_dataset.csv")
-        dataset_id="test_dataset",
+    rows = df.to_dict(orient="records")
-        rows_in_page=3,
-    )
-    assert len(rows.rows) == 3
 
-    scoring_fns_list = llama_stack_client.scoring_functions.list()
     scoring_functions = {
-        scoring_fns_list[0].identifier: None,
+        scoring_fn_id: None,
     }
 
     response = llama_stack_client.scoring.score(
-        input_rows=rows.rows,
+        input_rows=rows,
         scoring_functions=scoring_functions,
     )
     assert len(response.results) == len(scoring_functions)
     for x in scoring_functions:
         assert x in response.results
-        assert len(response.results[x].score_rows) == len(rows.rows)
+        assert len(response.results[x].score_rows) == len(rows)
-
-    # score batch
-    response = llama_stack_client.scoring.score_batch(
-        dataset_id="test_dataset",
-        scoring_functions=scoring_functions,
-        save_results_dataset=False,
-    )
-    assert len(response.results) == len(scoring_functions)
-    for x in scoring_functions:
-        assert x in response.results
-        assert len(response.results[x].score_rows) == 5
 
 
 def test_scoring_score_with_params_llm_as_judge(
-    llama_stack_client, sample_judge_prompt_template, judge_model_id, rag_dataset_for_test
+    llama_stack_client,
+    sample_judge_prompt_template,
+    judge_model_id,
 ):
     # scoring individual rows
-    rows = llama_stack_client.datasetio.get_rows_paginated(
+    df = pd.read_csv(Path(__file__).parent.parent / "datasets" / "test_dataset.csv")
-        dataset_id="test_dataset",
+    rows = df.to_dict(orient="records")
-        rows_in_page=3,
-    )
-    assert len(rows.rows) == 3
 
     scoring_functions = {
         "llm-as-judge::base": dict(
@@ -149,24 +122,13 @@ def test_scoring_score_with_params_llm_as_judge(
     }
 
     response = llama_stack_client.scoring.score(
-        input_rows=rows.rows,
+        input_rows=rows,
         scoring_functions=scoring_functions,
     )
     assert len(response.results) == len(scoring_functions)
     for x in scoring_functions:
         assert x in response.results
-        assert len(response.results[x].score_rows) == len(rows.rows)
+        assert len(response.results[x].score_rows) == len(rows)
-
-    # score batch
-    response = llama_stack_client.scoring.score_batch(
-        dataset_id="test_dataset",
-        scoring_functions=scoring_functions,
-        save_results_dataset=False,
-    )
-    assert len(response.results) == len(scoring_functions)
-    for x in scoring_functions:
-        assert x in response.results
-        assert len(response.results[x].score_rows) == 5
 
 
 @pytest.mark.parametrize(
@@ -178,13 +140,14 @@ def test_scoring_score_with_params_llm_as_judge(
     ],
 )
 def test_scoring_score_with_aggregation_functions(
-    llama_stack_client, sample_judge_prompt_template, judge_model_id, provider_id, rag_dataset_for_test
+    llama_stack_client,
+    sample_judge_prompt_template,
+    judge_model_id,
+    provider_id,
+    rag_dataset_for_test,
 ):
-    rows = llama_stack_client.datasetio.get_rows_paginated(
+    df = pd.read_csv(Path(__file__).parent.parent / "datasets" / "test_dataset.csv")
-        dataset_id="test_dataset",
+    rows = df.to_dict(orient="records")
-        rows_in_page=3,
-    )
-    assert len(rows.rows) == 3
 
     scoring_fns_list = [x for x in llama_stack_client.scoring_functions.list() if x.provider_id == provider_id]
     if len(scoring_fns_list) == 0:
@@ -224,12 +187,12 @@ def test_scoring_score_with_aggregation_functions(
         scoring_functions[scoring_fn.identifier] = None
 
     response = llama_stack_client.scoring.score(
-        input_rows=rows.rows,
+        input_rows=rows,
         scoring_functions=scoring_functions,
     )
 
     assert len(response.results) == len(scoring_functions)
     for x in scoring_functions:
         assert x in response.results
-        assert len(response.results[x].score_rows) == len(rows.rows)
+        assert len(response.results[x].score_rows) == len(rows)
         assert len(response.results[x].aggregated_results) == len(aggr_fns)

tests/unit/providers/inference/test_remote_vllm.py

@@ -187,8 +187,8 @@ def test_chat_completion_doesnt_block_event_loop(caplog):
     loop.set_debug(True)
     caplog.set_level(logging.WARNING)
 
-    # Log when event loop is blocked for more than 100ms
+    # Log when event loop is blocked for more than 200ms
-    loop.slow_callback_duration = 0.1
+    loop.slow_callback_duration = 0.2
     # Sleep for 500ms in our delayed http response
     sleep_time = 0.5