From 4cc1958af97a959a8d358d1008209c138067f362 Mon Sep 17 00:00:00 2001
From: Xi Yan <xiyan@meta.com>
Date: Wed, 12 Mar 2025 21:37:13 -0700
Subject: [PATCH] huggingface obey consistency

---
 docs/_static/llama-stack-spec.html    | 69 ++++++++++++++++-----------
 docs/_static/llama-stack-spec.yaml    | 40 +++++++++-------
 llama_stack/apis/datasets/datasets.py | 16 +++++--
 3 files changed, 76 insertions(+), 49 deletions(-)

diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html
index 47f48df3c..44459f2b9 100644
--- a/docs/_static/llama-stack-spec.html
+++ b/docs/_static/llama-stack-spec.html
@@ -6902,44 +6902,55 @@
                     "type": {
                         "type": "string",
                         "const": "huggingface",
-                        "default": "huggingface"
+                        "default": "huggingface",
+                        "description": "The type of the data source."
                     },
-                    "path": {
-                        "type": "string",
-                        "description": "The path to the dataset in Huggingface. E.g. - \"llamastack/simpleqa\""
-                    },
-                    "params": {
+                    "huggingface": {
                         "type": "object",
-                        "additionalProperties": {
-                            "oneOf": [
-                                {
-                                    "type": "null"
+                        "properties": {
+                            "path": {
+                                "type": "string",
+                                "description": "The path to the dataset in Huggingface. E.g. - \"llamastack/simpleqa\""
+                            },
+                            "params": {
+                                "type": "object",
+                                "additionalProperties": {
+                                    "oneOf": [
+                                        {
+                                            "type": "null"
+                                        },
+                                        {
+                                            "type": "boolean"
+                                        },
+                                        {
+                                            "type": "number"
+                                        },
+                                        {
+                                            "type": "string"
+                                        },
+                                        {
+                                            "type": "array"
+                                        },
+                                        {
+                                            "type": "object"
+                                        }
+                                    ]
                                 },
-                                {
-                                    "type": "boolean"
-                                },
-                                {
-                                    "type": "number"
-                                },
-                                {
-                                    "type": "string"
-                                },
-                                {
-                                    "type": "array"
-                                },
-                                {
-                                    "type": "object"
-                                }
-                            ]
+                                "description": "The parameters for the dataset."
+                            }
                         },
-                        "description": "The parameters for the dataset."
+                        "additionalProperties": false,
+                        "required": [
+                            "path",
+                            "params"
+                        ],
+                        "description": "The fields for a Huggingface dataset."
                     }
                 },
                 "additionalProperties": false,
                 "required": [
                     "type",
-                    "path",
-                    "params"
+                    "huggingface"
                 ],
                 "title": "HuggingfaceDataSource",
                 "description": "A dataset stored in Huggingface."
diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml
index 16ef6fed4..a257b2a7d 100644
--- a/docs/_static/llama-stack-spec.yaml
+++ b/docs/_static/llama-stack-spec.yaml
@@ -4775,26 +4775,34 @@ components:
           type: string
           const: huggingface
           default: huggingface
-        path:
-          type: string
-          description: >-
-            The path to the dataset in Huggingface. E.g. - "llamastack/simpleqa"
-        params:
+          description: The type of the data source.
+        huggingface:
           type: object
-          additionalProperties:
-            oneOf:
-              - type: 'null'
-              - type: boolean
-              - type: number
-              - type: string
-              - type: array
-              - type: object
-          description: The parameters for the dataset.
+          properties:
+            path:
+              type: string
+              description: >-
+                The path to the dataset in Huggingface. E.g. - "llamastack/simpleqa"
+            params:
+              type: object
+              additionalProperties:
+                oneOf:
+                  - type: 'null'
+                  - type: boolean
+                  - type: number
+                  - type: string
+                  - type: array
+                  - type: object
+              description: The parameters for the dataset.
+          additionalProperties: false
+          required:
+            - path
+            - params
+          description: The fields for a Huggingface dataset.
       additionalProperties: false
       required:
         - type
-        - path
-        - params
+        - huggingface
       title: HuggingfaceDataSource
       description: A dataset stored in Huggingface.
     RowsDataSource:
diff --git a/llama_stack/apis/datasets/datasets.py b/llama_stack/apis/datasets/datasets.py
index 36f75d7b3..20587a29e 100644
--- a/llama_stack/apis/datasets/datasets.py
+++ b/llama_stack/apis/datasets/datasets.py
@@ -55,18 +55,26 @@ class URIDataSource(BaseModel):
     uri: str
 
 
-@json_schema_type
-class HuggingfaceDataSource(BaseModel):
-    """A dataset stored in Huggingface.
+class HuggingfaceDataSourceFields(BaseModel):
+    """The fields for a Huggingface dataset.
     :param path: The path to the dataset in Huggingface. E.g.
         - "llamastack/simpleqa"
     :param params: The parameters for the dataset.
     """
-    type: Literal["huggingface"] = "huggingface"
     path: str
     params: Dict[str, Any]
 
 
+@json_schema_type
+class HuggingfaceDataSource(BaseModel):
+    """A dataset stored in Huggingface.
+    :param type: The type of the data source.
+    :param huggingface: The fields for a Huggingface dataset.
+    """
+    type: Literal["huggingface"] = "huggingface"
+    huggingface: HuggingfaceDataSourceFields
+
+
 @json_schema_type
 class RowsDataSource(BaseModel):
     """A dataset stored in rows.