Merge-related changes.

llama_stack/templates/template.py

@@ -11,11 +11,14 @@ 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.apis.preprocessing.preprocessors import PreprocessorInput
 from llama_stack.distribution.datatypes import (
     Api,
+    BenchmarkInput,
     BuildConfig,
+    DatasetInput,
     DistributionSpec,
     ModelInput,
     Provider,
@@ -29,7 +32,9 @@ from llama_stack.providers.utils.inference.model_registry import ProviderModelEn
 from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
 
 
-def get_model_registry(available_models: Dict[str, List[ProviderModelEntry]]) -> List[ModelInput]:
+def get_model_registry(
+    available_models: Dict[str, List[ProviderModelEntry]],
+) -> List[ModelInput]:
     models = []
     for provider_id, entries in available_models.items():
         for entry in entries:
@@ -57,6 +62,8 @@ class RunConfigSettings(BaseModel):
     default_models: Optional[List[ModelInput]] = None
     default_shields: Optional[List[ShieldInput]] = None
     default_tool_groups: Optional[List[ToolGroupInput]] = None
+    default_datasets: Optional[List[DatasetInput]] = None
+    default_benchmarks: Optional[List[BenchmarkInput]] = None
     default_preprocessors: Optional[List[PreprocessorInput]] = None
 
     def run_config(
@@ -115,6 +122,8 @@ class RunConfigSettings(BaseModel):
             models=self.default_models or [],
             shields=self.default_shields or [],
             tool_groups=self.default_tool_groups or [],
+            datasets=self.default_datasets or [],
+            benchmarks=self.default_benchmarks or [],
             preprocessors=self.default_preprocessors or [],
         )
 
@@ -190,7 +199,7 @@ class DistributionTemplate(BaseModel):
                     default_models.append(
                         DefaultModel(
                             model_id=model_entry.provider_model_id,
-                            doc_string=f"({' -- '.join(doc_parts)})" if doc_parts else "",
+                            doc_string=(f"({' -- '.join(doc_parts)})" if doc_parts else ""),
                         )
                     )
 
@@ -209,7 +218,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)