feat(api): define a more coherent jobs api across different flows

Signed-off-by: Ihar Hrachyshka <ihar.hrachyshka@gmail.com>
This commit is contained in:
Ihar Hrachyshka 2025-03-24 20:54:04 -04:00
parent 71ed47ea76
commit 0f50cfa561
15 changed files with 1864 additions and 1670 deletions

View file

@ -5,10 +5,11 @@
# the root directory of this source tree.
from enum import Enum
from typing import Any, Dict, List, Optional, Protocol, Union
from typing import List, Literal, Optional, Protocol
from pydantic import BaseModel
from llama_stack.apis.common.job_types import BaseJob
from llama_stack.apis.inference import Message
from llama_stack.schema_utils import json_schema_type, webmethod
@ -34,11 +35,13 @@ class SyntheticDataGenerationRequest(BaseModel):
@json_schema_type
class SyntheticDataGenerationResponse(BaseModel):
"""Response from the synthetic data generation. Batch of (prompt, response, score) tuples that pass the threshold."""
class SyntheticDataGenerationJob(BaseJob, BaseModel):
type: Literal["synthetic-data-generation"] = "synthetic-data-generation"
synthetic_data: List[Dict[str, Any]]
statistics: Optional[Dict[str, Any]] = None
@json_schema_type
class ListSyntheticDataGenerationJobsResponse(BaseModel):
items: list[SyntheticDataGenerationJob]
class SyntheticDataGeneration(Protocol):
@ -48,4 +51,24 @@ class SyntheticDataGeneration(Protocol):
dialogs: List[Message],
filtering_function: FilteringFunction = FilteringFunction.none,
model: Optional[str] = None,
) -> Union[SyntheticDataGenerationResponse]: ...
) -> SyntheticDataGenerationJob: ...
# CRUD operations on running jobs
@webmethod(route="/synthetic-data-generation/jobs/{job_id:path}", method="GET")
async def get_synthetic_data_generation_job(self) -> SyntheticDataGenerationJob: ...
@webmethod(route="/synthetic-data-generation/jobs", method="GET")
async def list_synthetic_data_generation_jobs(self) -> ListSyntheticDataGenerationJobsResponse: ...
@webmethod(route="/synthetic-data-generation/jobs/{job_id:path}", method="POST")
async def update_synthetic_data_generation_job(
self, job: SyntheticDataGenerationJob
) -> SyntheticDataGenerationJob: ...
@webmethod(route="/synthetic-data-generation/job/{job_id:path}", method="DELETE")
async def delete_synthetic_data_generation_job(self, job_id: str) -> None: ...
# Note: pause/resume/cancel are achieved as follows:
# - POST with status=paused
# - POST with status=resuming
# - POST with status=cancelled