move jobs, fix errors

llama_stack/apis/common/job_types.py

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+from llama_models.schema_utils import json_schema_type
+from pydantic import BaseModel
+
+
+@json_schema_type
+class Job(BaseModel):
+    job_id: str

llama_stack/apis/datasets/datasets.py

@@ -15,17 +15,12 @@ from pydantic import BaseModel, Field
 from llama_stack.apis.common.type_system import ParamType
 
 
-@json_schema_type
-class DatasetSchema(BaseModel):
-    columns: Dict[str, ParamType]
-
-
 @json_schema_type
 class DatasetDef(BaseModel):
     identifier: str = Field(
         description="A unique name for the dataset",
     )
-    dataset_schema: DatasetSchema = Field(
+    columns_schema: Dict[str, ParamType] = Field(
         description="The schema definition for this dataset",
     )
     url: URL

llama_stack/apis/eval/eval.py

@@ -11,6 +11,9 @@ from typing_extensions import Annotated
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_models.schema_utils import json_schema_type, webmethod
 from llama_stack.apis.scoring_functions import *  # noqa: F403
+from llama_stack.apis.agents import AgentConfig
+from llama_stack.apis.common.job_types import Job
+from llama_stack.apis.scoring import *  # noqa: F403
 
 
 @json_schema_type
@@ -32,11 +35,6 @@ EvalCandidate = Annotated[
 ]
 
 
-@json_schema_type
-class Job(BaseModel):
-    job_id: str
-
-
 @json_schema_type
 class EvaluateResponse(BaseModel):
     generations: List[Dict[str, Any]]

llama_stack/apis/scoring_functions/scoring_functions.py

@@ -22,6 +22,7 @@ from typing_extensions import Annotated
 from llama_stack.apis.common.type_system import ParamType
 
 
+@json_schema_type
 class Parameter(BaseModel):
     name: str
     type: ParamType
@@ -32,6 +33,7 @@ class Parameter(BaseModel):
 # with standard metrics so they can be rolled up?
 
 
+@json_schema_type
 class CommonDef(BaseModel):
     name: str
     description: Optional[str] = None
@@ -39,8 +41,11 @@ class CommonDef(BaseModel):
         default_factory=dict,
         description="Any additional metadata for this definition",
     )
+    # Hack: same with memory_banks for union defs
+    provider_id: str = ""
 
 
+@json_schema_type
 class DeterministicFunctionDef(CommonDef):
     type: Literal["deterministic"] = "deterministic"
     parameters: List[Parameter] = Field(
@@ -52,6 +57,7 @@ class DeterministicFunctionDef(CommonDef):
     # We can optionally add information here to support packaging of code, etc.
 
 
+@json_schema_type
 class LLMJudgeFunctionDef(CommonDef):
     type: Literal["judge"] = "judge"
     model: str = Field(
@@ -63,12 +69,7 @@ ScoringFunctionDef = Annotated[
     Union[DeterministicFunctionDef, LLMJudgeFunctionDef], Field(discriminator="type")
 ]
 
-
-@json_schema_type
-class ScoringFunctionDefWithProvider(ScoringFunctionDef):
-    provider_id: str = Field(
-        description="The provider ID for this scoring function",
-    )
+ScoringFunctionDefWithProvider = ScoringFunctionDef
 
 
 @runtime_checkable