generate openapi

llama_stack/apis/common/type_system.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Dict, List, Literal, Union
+from typing import Literal, Union
 
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated
@@ -24,12 +24,10 @@ class BooleanType(BaseModel):
 
 class ArrayType(BaseModel):
     type: Literal["array"] = "array"
-    items: "ParamType"
 
 
 class ObjectType(BaseModel):
     type: Literal["object"] = "object"
-    properties: Dict[str, "ParamType"] = Field(default_factory=dict)
 
 
 class JsonType(BaseModel):
@@ -38,7 +36,6 @@ class JsonType(BaseModel):
 
 class UnionType(BaseModel):
     type: Literal["union"] = "union"
-    options: List["ParamType"] = Field(default_factory=list)
 
 
 class CustomType(BaseModel):
@@ -77,7 +74,3 @@ ParamType = Annotated[
     ],
     Field(discriminator="type"),
 ]
-
-ArrayType.model_rebuild()
-ObjectType.model_rebuild()
-UnionType.model_rebuild()

llama_stack/apis/post_training/post_training.py

@@ -14,7 +14,7 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel, Field
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.apis.dataset import *  # noqa: F403
+from llama_stack.apis.datasets import *  # noqa: F403
 from llama_stack.apis.common.training_types import *  # noqa: F403
 
 
@@ -107,8 +107,8 @@ class PostTrainingSFTRequest(BaseModel):
     job_uuid: str
 
     model: str
-    dataset: TrainEvalDataset
-    validation_dataset: TrainEvalDataset
+    dataset: str
+    validation_dataset: str
 
     algorithm: FinetuningAlgorithm
     algorithm_config: Union[
@@ -131,8 +131,8 @@ class PostTrainingRLHFRequest(BaseModel):
 
     finetuned_model: URL
 
-    dataset: TrainEvalDataset
-    validation_dataset: TrainEvalDataset
+    dataset: str
+    validation_dataset: str
 
     algorithm: RLHFAlgorithm
     algorithm_config: Union[DPOAlignmentConfig]
@@ -181,8 +181,8 @@ class PostTraining(Protocol):
         self,
         job_uuid: str,
         model: str,
-        dataset: TrainEvalDataset,
-        validation_dataset: TrainEvalDataset,
+        dataset: str,
+        validation_dataset: str,
         algorithm: FinetuningAlgorithm,
         algorithm_config: Union[
             LoraFinetuningConfig, QLoraFinetuningConfig, DoraFinetuningConfig
@@ -198,8 +198,8 @@ class PostTraining(Protocol):
         self,
         job_uuid: str,
         finetuned_model: URL,
-        dataset: TrainEvalDataset,
-        validation_dataset: TrainEvalDataset,
+        dataset: str,
+        validation_dataset: str,
         algorithm: RLHFAlgorithm,
         algorithm_config: Union[DPOAlignmentConfig],
         optimizer_config: OptimizerConfig,

llama_stack/apis/synthetic_data_generation/synthetic_data_generation.py

@@ -13,7 +13,6 @@ from llama_models.schema_utils import json_schema_type, webmethod
 from pydantic import BaseModel
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
-from llama_stack.apis.reward_scoring import *  # noqa: F403
 
 
 class FilteringFunction(Enum):
@@ -40,7 +39,7 @@ class SyntheticDataGenerationRequest(BaseModel):
 class SyntheticDataGenerationResponse(BaseModel):
     """Response from the synthetic data generation. Batch of (prompt, response, score) tuples that pass the threshold."""
 
-    synthetic_data: List[ScoredDialogGenerations]
+    synthetic_data: List[Dict[str, Any]]
     statistics: Optional[Dict[str, Any]] = None