persist registered objects with distribution

llama_stack/apis/datasets/datasets.py

@@ -10,10 +10,10 @@ from llama_models.llama3.api.datatypes import URL
 
 from llama_models.schema_utils import json_schema_type, webmethod
 
-from pydantic import BaseModel, Field
-
 from llama_stack.apis.common.type_system import ParamType
 
+from pydantic import BaseModel, Field
+
 
 @json_schema_type
 class DatasetDef(BaseModel):
@@ -24,7 +24,7 @@ class DatasetDef(BaseModel):
         description="The schema definition for this dataset",
     )
     url: URL
-    metadata: Dict[str, Any] = Field(
+    metadata: Dict[str, str] = Field(
         default_factory=dict,
         description="Any additional metadata for this dataset",
     )

llama_stack/apis/models/models.py

@@ -17,7 +17,7 @@ class ModelDef(BaseModel):
     llama_model: str = Field(
         description="Pointer to the underlying core Llama family model. Each model served by Llama Stack must have a core Llama model.",
     )
-    metadata: Dict[str, Any] = Field(
+    metadata: Dict[str, str] = Field(
         default_factory=dict,
         description="Any additional metadata for this model",
     )

llama_stack/apis/scoring_functions/scoring_functions.py

@@ -4,12 +4,12 @@
 # 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, Protocol, runtime_checkable
+from typing import Dict, List, Optional, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
-from pydantic import BaseModel, Field
 
-from llama_stack.apis.common.type_system import ParamType
+from llama_stack.apis.common.type_system import JsonType, ParamType
+from pydantic import BaseModel, Field
 
 
 @json_schema_type
@@ -36,7 +36,7 @@ class LLMAsJudgeContext(BaseModel):
 class ScoringFnDef(BaseModel):
     identifier: str
     description: Optional[str] = None
-    metadata: Dict[str, Any] = Field(
+    metadata: Dict[str, str] = Field(
         default_factory=dict,
         description="Any additional metadata for this definition",
     )

llama_stack/apis/shields/shields.py

@@ -26,7 +26,7 @@ class ShieldDef(BaseModel):
     type: str = Field(
         description="The type of shield this is; the value is one of the ShieldType enum"
     )
-    params: Dict[str, Any] = Field(
+    params: Dict[str, str] = Field(
         default_factory=dict,
         description="Any additional parameters needed for this shield",
     )