Merge branch 'api_2' into api_3

llama_stack/apis/datasets/datasets.py

@@ -16,7 +16,7 @@ from llama_stack.schema_utils import json_schema_type, register_schema, webmetho
 class Schema(Enum):
     """
     Schema of the dataset. Each type has a different column format.
-    :cvar jsonl_messages: The dataset is a JSONL file with messages. Examples:
+    :cvar messages: The dataset contains messages used for post-training. Examples:
         {
             "messages": [
                 {"role": "user", "content": "Hello, world!"},
@@ -25,7 +25,7 @@ class Schema(Enum):
         }
     """
 
-    jsonl_messages = "jsonl_messages"
+    messages = "messages"
     # TODO: add more schemas here
 
 
@@ -36,36 +36,36 @@ class DatasetType(Enum):
 
 
 @json_schema_type
-class URIDataReference(BaseModel):
+class URIDataSource(BaseModel):
     type: Literal["uri"] = "uri"
     uri: str
 
 
 @json_schema_type
-class HuggingfaceDataReference(BaseModel):
+class HuggingfaceDataSource(BaseModel):
     type: Literal["huggingface"] = "huggingface"
     dataset_path: str
     params: Dict[str, Any]
 
 
 @json_schema_type
-class RowsDataReference(BaseModel):
+class RowsDataSource(BaseModel):
     type: Literal["rows"] = "rows"
     rows: List[Dict[str, Any]]
 
 
-DataReference = register_schema(
+DataSource = register_schema(
     Annotated[
-        Union[URIDataReference, HuggingfaceDataReference, RowsDataReference],
+        Union[URIDataSource, HuggingfaceDataSource, RowsDataSource],
         Field(discriminator="type"),
     ],
-    name="DataReference",
+    name="DataSource",
 )
 
 
 class CommonDatasetFields(BaseModel):
     schema: Schema
-    data_reference: DataReference
+    data_source: DataSource
     metadata: Dict[str, Any] = Field(
         default_factory=dict,
         description="Any additional metadata for this dataset",
@@ -100,16 +100,16 @@ class Datasets(Protocol):
     async def register_dataset(
         self,
         schema: Schema,
-        data_reference: DataReference,
+        data_source: DataSource,
         metadata: Optional[Dict[str, Any]] = None,
         dataset_id: Optional[str] = None,
     ) -> Dataset:
         """
-        Register a new dataset through a file or
+        Register a new dataset.
 
         :param schema: The schema format of the dataset. One of
-            - jsonl_messages: The dataset is a JSONL file with messages in column format
-        :param data_reference: The data reference of the dataset. Examples:
+            - messages: The dataset contains a messages column with list of messages for post-training.
+        :param data_source: The data source of the dataset. Examples:
            - {
                "type": "uri",
                "uri": "https://mywebsite.com/mydata.jsonl"

llama_stack/apis/scoring_functions/scoring_functions.py

@@ -12,8 +12,8 @@ from typing import (
     Literal,
     Optional,
     Protocol,
-    Union,
     runtime_checkable,
+    Union,
 )
 
 from pydantic import BaseModel, Field
@@ -218,7 +218,9 @@ class CommonScoringFnFields(BaseModel):
 
 @json_schema_type
 class ScoringFn(CommonScoringFnFields, Resource):
-    type: Literal[ResourceType.scoring_function.value] = ResourceType.scoring_function.value
+    type: Literal[ResourceType.scoring_function.value] = (
+        ResourceType.scoring_function.value
+    )
 
     @property
     def scoring_fn_id(self) -> str:
@@ -245,13 +247,14 @@ class ScoringFunctions(Protocol):
     async def list_scoring_functions(self) -> ListScoringFunctionsResponse: ...
 
     @webmethod(route="/scoring-functions/{scoring_fn_id:path}", method="GET")
-    async def get_scoring_function(self, scoring_fn_id: str, /) -> Optional[ScoringFn]: ...
+    async def get_scoring_function(
+        self, scoring_fn_id: str, /
+    ) -> Optional[ScoringFn]: ...
 
     @webmethod(route="/scoring-functions", method="POST")
     async def register_scoring_function(
         self,
-        scoring_fn_type: ScoringFunctionType,
-        params: ScoringFnParams = None,
+        fn: ScoringFnParams,
         scoring_fn_id: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
     ) -> ScoringFn:
@@ -259,8 +262,7 @@ class ScoringFunctions(Protocol):
         Register a new scoring function with given parameters.
         Only valid scoring function type that can be parameterized can be registered.
 
-        :param scoring_fn_type: The type of scoring function to register.
-        :param params: The parameters for the scoring function.
+        :param fn: The type and parameters for the scoring function.
         :param scoring_fn_id: (Optional) The ID of the scoring function to register. If not provided, a random ID will be generated.
         :param metadata: (Optional) Any additional metadata to be associated with the scoring function.
             - E.g. {"description": "This scoring function is used for ..."}