feat(api): clean up api for evals

- as title
This commit is contained in:
Xi Yan 2025-03-11 14:17:08 -07:00
parent 5f90be5388
commit 62839c66a1

View file

@ -4,19 +4,35 @@
# This source code is licensed under the terms described in the LICENSE file in # This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree. # the root directory of this source tree.
from enum import Enum
from typing import Any, Dict, List, Literal, Optional, Protocol from typing import Any, Dict, List, Literal, Optional, Protocol
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from llama_stack.apis.common.content_types import URL
from llama_stack.apis.common.type_system import ParamType
from llama_stack.apis.resource import Resource, ResourceType from llama_stack.apis.resource import Resource, ResourceType
from llama_stack.schema_utils import json_schema_type, webmethod from llama_stack.schema_utils import json_schema_type, webmethod
@json_schema_type
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:
{
"messages": [
{"role": "user", "content": "Hello, world!"},
{"role": "assistant", "content": "Hello, world!"},
]
}
"""
jsonl_messages = "jsonl_messages"
class CommonDatasetFields(BaseModel): class CommonDatasetFields(BaseModel):
dataset_schema: Dict[str, ParamType] schema: Schema
url: URL uri: str
metadata: Dict[str, Any] = Field( metadata: Dict[str, Any] = Field(
default_factory=dict, default_factory=dict,
description="Any additional metadata for this dataset", description="Any additional metadata for this dataset",
@ -50,13 +66,10 @@ class Datasets(Protocol):
@webmethod(route="/datasets", method="POST") @webmethod(route="/datasets", method="POST")
async def register_dataset( async def register_dataset(
self, self,
dataset_id: str, schema: Schema,
dataset_schema: Dict[str, ParamType], uri: str,
url: URL,
provider_dataset_id: Optional[str] = None,
provider_id: Optional[str] = None,
metadata: Optional[Dict[str, Any]] = None, metadata: Optional[Dict[str, Any]] = None,
) -> None: ... ) -> Dataset: ...
@webmethod(route="/datasets/{dataset_id:path}", method="GET") @webmethod(route="/datasets/{dataset_id:path}", method="GET")
async def get_dataset( async def get_dataset(