llama-stack-api 0.4.4__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

llama_stack_api/post_training/fastapi_routes.py

@@ -0,0 +1,116 @@
+# 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.
+
+"""FastAPI router for the Post-Training API.
+
+This module defines the FastAPI router for the Post-Training API using standard
+FastAPI route decorators.
+"""
+
+from typing import Annotated
+
+from fastapi import APIRouter, Body, Depends
+
+from llama_stack_api.router_utils import create_path_dependency, standard_responses
+from llama_stack_api.version import LLAMA_STACK_API_V1ALPHA
+
+from .api import PostTraining
+from .models import (
+    CancelTrainingJobRequest,
+    GetTrainingJobArtifactsRequest,
+    GetTrainingJobStatusRequest,
+    ListPostTrainingJobsResponse,
+    PostTrainingJob,
+    PostTrainingJobArtifactsResponse,
+    PostTrainingJobStatusResponse,
+    PreferenceOptimizeRequest,
+    SupervisedFineTuneRequest,
+)
+
+# Path parameter dependencies for single-field models
+get_training_job_status_request = create_path_dependency(GetTrainingJobStatusRequest)
+cancel_training_job_request = create_path_dependency(CancelTrainingJobRequest)
+get_training_job_artifacts_request = create_path_dependency(GetTrainingJobArtifactsRequest)
+
+
+def create_router(impl: PostTraining) -> APIRouter:
+    """Create a FastAPI router for the Post-Training API."""
+    router = APIRouter(
+        prefix=f"/{LLAMA_STACK_API_V1ALPHA}",
+        tags=["Post Training"],
+        responses=standard_responses,
+    )
+
+    @router.post(
+        "/post-training/supervised-fine-tune",
+        response_model=PostTrainingJob,
+        summary="Run supervised fine-tuning of a model.",
+        description="Run supervised fine-tuning of a model.",
+        responses={200: {"description": "A PostTrainingJob."}},
+    )
+    async def supervised_fine_tune(
+        request: Annotated[SupervisedFineTuneRequest, Body(...)],
+    ) -> PostTrainingJob:
+        return await impl.supervised_fine_tune(request)
+
+    @router.post(
+        "/post-training/preference-optimize",
+        response_model=PostTrainingJob,
+        summary="Run preference optimization of a model.",
+        description="Run preference optimization of a model.",
+        responses={200: {"description": "A PostTrainingJob."}},
+    )
+    async def preference_optimize(
+        request: Annotated[PreferenceOptimizeRequest, Body(...)],
+    ) -> PostTrainingJob:
+        return await impl.preference_optimize(request)
+
+    @router.get(
+        "/post-training/jobs",
+        response_model=ListPostTrainingJobsResponse,
+        summary="Get all training jobs.",
+        description="Get all training jobs.",
+        responses={200: {"description": "A ListPostTrainingJobsResponse."}},
+    )
+    async def get_training_jobs() -> ListPostTrainingJobsResponse:
+        return await impl.get_training_jobs()
+
+    @router.get(
+        "/post-training/job/status",
+        response_model=PostTrainingJobStatusResponse,
+        summary="Get the status of a training job.",
+        description="Get the status of a training job.",
+        responses={200: {"description": "A PostTrainingJobStatusResponse."}},
+    )
+    async def get_training_job_status(
+        request: Annotated[GetTrainingJobStatusRequest, Depends(get_training_job_status_request)],
+    ) -> PostTrainingJobStatusResponse:
+        return await impl.get_training_job_status(request)
+
+    @router.post(
+        "/post-training/job/cancel",
+        summary="Cancel a training job.",
+        description="Cancel a training job.",
+        responses={200: {"description": "Successfully cancelled the training job."}},
+    )
+    async def cancel_training_job(
+        request: Annotated[CancelTrainingJobRequest, Depends(cancel_training_job_request)],
+    ) -> None:
+        return await impl.cancel_training_job(request)
+
+    @router.get(
+        "/post-training/job/artifacts",
+        response_model=PostTrainingJobArtifactsResponse,
+        summary="Get the artifacts of a training job.",
+        description="Get the artifacts of a training job.",
+        responses={200: {"description": "A PostTrainingJobArtifactsResponse."}},
+    )
+    async def get_training_job_artifacts(
+        request: Annotated[GetTrainingJobArtifactsRequest, Depends(get_training_job_artifacts_request)],
+    ) -> PostTrainingJobArtifactsResponse:
+        return await impl.get_training_job_artifacts(request)
+
+    return router

llama_stack_api/{post_training.py → post_training/models.py}

@@ -4,17 +4,22 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+"""Pydantic models for Post-Training API requests and responses.
+
+This module defines the request and response models for the Post-Training API
+using Pydantic with Field descriptions for OpenAPI schema generation.
+"""
+
 from datetime import datetime
 from enum import Enum
-from typing import Annotated, Any, Literal, Protocol
+from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, Field
 
 from llama_stack_api.common.content_types import URL
 from llama_stack_api.common.job_types import JobStatus
 from llama_stack_api.common.training_types import Checkpoint
-from llama_stack_api.schema_utils import json_schema_type, register_schema, webmethod
-from llama_stack_api.version import LLAMA_STACK_API_V1ALPHA
+from llama_stack_api.schema_utils import json_schema_type, register_schema
 
 
 @json_schema_type
@@ -285,86 +290,50 @@ class PostTrainingJobArtifactsResponse(BaseModel):
     # TODO(ashwin): metrics, evals
 
 
-class PostTraining(Protocol):
-    @webmethod(route="/post-training/supervised-fine-tune", method="POST", level=LLAMA_STACK_API_V1ALPHA)
-    async def supervised_fine_tune(
-        self,
-        job_uuid: str,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
-        model: str | None = Field(
-            default=None,
-            description="Model descriptor for training if not in provider config`",
-        ),
-        checkpoint_dir: str | None = None,
-        algorithm_config: AlgorithmConfig | None = None,
-    ) -> PostTrainingJob:
-        """Run supervised fine-tuning of a model.
-
-        :param job_uuid: The UUID of the job to create.
-        :param training_config: The training configuration.
-        :param hyperparam_search_config: The hyperparam search configuration.
-        :param logger_config: The logger configuration.
-        :param model: The model to fine-tune.
-        :param checkpoint_dir: The directory to save checkpoint(s) to.
-        :param algorithm_config: The algorithm configuration.
-        :returns: A PostTrainingJob.
-        """
-        ...
-
-    @webmethod(route="/post-training/preference-optimize", method="POST", level=LLAMA_STACK_API_V1ALPHA)
-    async def preference_optimize(
-        self,
-        job_uuid: str,
-        finetuned_model: str,
-        algorithm_config: DPOAlignmentConfig,
-        training_config: TrainingConfig,
-        hyperparam_search_config: dict[str, Any],
-        logger_config: dict[str, Any],
-    ) -> PostTrainingJob:
-        """Run preference optimization of a model.
-
-        :param job_uuid: The UUID of the job to create.
-        :param finetuned_model: The model to fine-tune.
-        :param algorithm_config: The algorithm configuration.
-        :param training_config: The training configuration.
-        :param hyperparam_search_config: The hyperparam search configuration.
-        :param logger_config: The logger configuration.
-        :returns: A PostTrainingJob.
-        """
-        ...
-
-    @webmethod(route="/post-training/jobs", method="GET", level=LLAMA_STACK_API_V1ALPHA)
-    async def get_training_jobs(self) -> ListPostTrainingJobsResponse:
-        """Get all training jobs.
-
-        :returns: A ListPostTrainingJobsResponse.
-        """
-        ...
-
-    @webmethod(route="/post-training/job/status", method="GET", level=LLAMA_STACK_API_V1ALPHA)
-    async def get_training_job_status(self, job_uuid: str) -> PostTrainingJobStatusResponse:
-        """Get the status of a training job.
-
-        :param job_uuid: The UUID of the job to get the status of.
-        :returns: A PostTrainingJobStatusResponse.
-        """
-        ...
-
-    @webmethod(route="/post-training/job/cancel", method="POST", level=LLAMA_STACK_API_V1ALPHA)
-    async def cancel_training_job(self, job_uuid: str) -> None:
-        """Cancel a training job.
-
-        :param job_uuid: The UUID of the job to cancel.
-        """
-        ...
-
-    @webmethod(route="/post-training/job/artifacts", method="GET", level=LLAMA_STACK_API_V1ALPHA)
-    async def get_training_job_artifacts(self, job_uuid: str) -> PostTrainingJobArtifactsResponse:
-        """Get the artifacts of a training job.
-
-        :param job_uuid: The UUID of the job to get the artifacts of.
-        :returns: A PostTrainingJobArtifactsResponse.
-        """
-        ...
+@json_schema_type
+class SupervisedFineTuneRequest(BaseModel):
+    """Request to run supervised fine-tuning of a model."""
+
+    job_uuid: str = Field(..., description="The UUID of the job to create.")
+    training_config: TrainingConfig = Field(..., description="The training configuration.")
+    hyperparam_search_config: dict[str, Any] = Field(..., description="The hyperparam search configuration.")
+    logger_config: dict[str, Any] = Field(..., description="The logger configuration.")
+    model: str | None = Field(
+        default=None,
+        description="Model descriptor for training if not in provider config",
+    )
+    checkpoint_dir: str | None = Field(default=None, description="The directory to save checkpoint(s) to.")
+    algorithm_config: AlgorithmConfig | None = Field(default=None, description="The algorithm configuration.")
+
+
+@json_schema_type
+class PreferenceOptimizeRequest(BaseModel):
+    """Request to run preference optimization of a model."""
+
+    job_uuid: str = Field(..., description="The UUID of the job to create.")
+    finetuned_model: str = Field(..., description="The model to fine-tune.")
+    algorithm_config: DPOAlignmentConfig = Field(..., description="The algorithm configuration.")
+    training_config: TrainingConfig = Field(..., description="The training configuration.")
+    hyperparam_search_config: dict[str, Any] = Field(..., description="The hyperparam search configuration.")
+    logger_config: dict[str, Any] = Field(..., description="The logger configuration.")
+
+
+@json_schema_type
+class GetTrainingJobStatusRequest(BaseModel):
+    """Request to get the status of a training job."""
+
+    job_uuid: str = Field(..., description="The UUID of the job to get the status of.")
+
+
+@json_schema_type
+class CancelTrainingJobRequest(BaseModel):
+    """Request to cancel a training job."""
+
+    job_uuid: str = Field(..., description="The UUID of the job to cancel.")
+
+
+@json_schema_type
+class GetTrainingJobArtifactsRequest(BaseModel):
+    """Request to get the artifacts of a training job."""
+
+    job_uuid: str = Field(..., description="The UUID of the job to get the artifacts of.")

llama_stack_api/prompts/__init__.py

@@ -0,0 +1,47 @@
+# 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.
+
+"""Prompts API protocol and models.
+
+This module contains the Prompts protocol definition.
+Pydantic models are defined in llama_stack_api.prompts.models.
+The FastAPI router is defined in llama_stack_api.prompts.fastapi_routes.
+"""
+
+# Import fastapi_routes for router factory access
+from . import fastapi_routes
+
+# Import protocol for FastAPI router
+from .api import Prompts
+
+# Import models for re-export
+from .models import (
+    CreatePromptRequest,
+    DeletePromptRequest,
+    GetPromptRequest,
+    ListPromptsResponse,
+    ListPromptVersionsRequest,
+    Prompt,
+    SetDefaultVersionBodyRequest,
+    SetDefaultVersionRequest,
+    UpdatePromptBodyRequest,
+    UpdatePromptRequest,
+)
+
+__all__ = [
+    "CreatePromptRequest",
+    "DeletePromptRequest",
+    "GetPromptRequest",
+    "ListPromptVersionsRequest",
+    "ListPromptsResponse",
+    "Prompt",
+    "Prompts",
+    "SetDefaultVersionBodyRequest",
+    "SetDefaultVersionRequest",
+    "UpdatePromptBodyRequest",
+    "UpdatePromptRequest",
+    "fastapi_routes",
+]

llama_stack_api/prompts/api.py

@@ -0,0 +1,44 @@
+# 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.
+
+"""Prompts API protocol definition.
+
+This module contains the Prompts protocol definition.
+Pydantic models are defined in llama_stack_api.prompts.models.
+The FastAPI router is defined in llama_stack_api.prompts.fastapi_routes.
+"""
+
+from typing import Protocol, runtime_checkable
+
+from .models import (
+    CreatePromptRequest,
+    DeletePromptRequest,
+    GetPromptRequest,
+    ListPromptsResponse,
+    ListPromptVersionsRequest,
+    Prompt,
+    SetDefaultVersionRequest,
+    UpdatePromptRequest,
+)
+
+
+@runtime_checkable
+class Prompts(Protocol):
+    """Protocol for prompt management operations."""
+
+    async def list_prompts(self) -> ListPromptsResponse: ...
+
+    async def list_prompt_versions(self, request: ListPromptVersionsRequest) -> ListPromptsResponse: ...
+
+    async def get_prompt(self, request: GetPromptRequest) -> Prompt: ...
+
+    async def create_prompt(self, request: CreatePromptRequest) -> Prompt: ...
+
+    async def update_prompt(self, request: UpdatePromptRequest) -> Prompt: ...
+
+    async def delete_prompt(self, request: DeletePromptRequest) -> None: ...
+
+    async def set_default_version(self, request: SetDefaultVersionRequest) -> Prompt: ...

llama_stack_api/prompts/fastapi_routes.py

@@ -0,0 +1,163 @@
+# 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.
+
+"""FastAPI router for the Prompts API.
+
+This module defines the FastAPI router for the Prompts API using standard
+FastAPI route decorators.
+"""
+
+from typing import Annotated
+
+from fastapi import APIRouter, Body, Depends, Path, Query
+
+from llama_stack_api.router_utils import create_path_dependency, standard_responses
+from llama_stack_api.version import LLAMA_STACK_API_V1
+
+from .api import Prompts
+from .models import (
+    CreatePromptRequest,
+    DeletePromptRequest,
+    GetPromptRequest,
+    ListPromptsResponse,
+    ListPromptVersionsRequest,
+    Prompt,
+    SetDefaultVersionBodyRequest,
+    SetDefaultVersionRequest,
+    UpdatePromptBodyRequest,
+    UpdatePromptRequest,
+)
+
+# Path parameter dependencies for single-field models
+list_prompt_versions_request = create_path_dependency(ListPromptVersionsRequest)
+delete_prompt_request = create_path_dependency(DeletePromptRequest)
+
+
+def create_router(impl: Prompts) -> APIRouter:
+    """Create a FastAPI router for the Prompts API.
+
+    Args:
+        impl: The Prompts implementation instance
+
+    Returns:
+        APIRouter configured for the Prompts API
+    """
+    router = APIRouter(
+        prefix=f"/{LLAMA_STACK_API_V1}",
+        tags=["Prompts"],
+        responses=standard_responses,
+    )
+
+    @router.get(
+        "/prompts",
+        response_model=ListPromptsResponse,
+        summary="List all prompts.",
+        description="List all prompts.",
+        responses={
+            200: {"description": "A ListPromptsResponse containing all prompts."},
+        },
+    )
+    async def list_prompts() -> ListPromptsResponse:
+        return await impl.list_prompts()
+
+    @router.get(
+        "/prompts/{prompt_id}/versions",
+        response_model=ListPromptsResponse,
+        summary="List prompt versions.",
+        description="List all versions of a specific prompt.",
+        responses={
+            200: {"description": "A ListPromptsResponse containing all versions of the prompt."},
+        },
+    )
+    async def list_prompt_versions(
+        request: Annotated[ListPromptVersionsRequest, Depends(list_prompt_versions_request)],
+    ) -> ListPromptsResponse:
+        return await impl.list_prompt_versions(request)
+
+    @router.get(
+        "/prompts/{prompt_id}",
+        response_model=Prompt,
+        summary="Get a prompt.",
+        description="Get a prompt by its identifier and optional version.",
+        responses={
+            200: {"description": "A Prompt resource."},
+        },
+    )
+    async def get_prompt(
+        prompt_id: Annotated[str, Path(description="The identifier of the prompt to get.")],
+        version: Annotated[
+            int | None, Query(description="The version of the prompt to get (defaults to latest).")
+        ] = None,
+    ) -> Prompt:
+        request = GetPromptRequest(prompt_id=prompt_id, version=version)
+        return await impl.get_prompt(request)
+
+    @router.post(
+        "/prompts",
+        response_model=Prompt,
+        summary="Create a prompt.",
+        description="Create a new prompt.",
+        responses={
+            200: {"description": "The created Prompt resource."},
+        },
+    )
+    async def create_prompt(
+        request: Annotated[CreatePromptRequest, Body(...)],
+    ) -> Prompt:
+        return await impl.create_prompt(request)
+
+    @router.put(
+        "/prompts/{prompt_id}",
+        response_model=Prompt,
+        summary="Update a prompt.",
+        description="Update an existing prompt (increments version).",
+        responses={
+            200: {"description": "The updated Prompt resource with incremented version."},
+        },
+    )
+    async def update_prompt(
+        prompt_id: Annotated[str, Path(description="The identifier of the prompt to update.")],
+        body: Annotated[UpdatePromptBodyRequest, Body(...)],
+    ) -> Prompt:
+        request = UpdatePromptRequest(
+            prompt_id=prompt_id,
+            prompt=body.prompt,
+            version=body.version,
+            variables=body.variables,
+            set_as_default=body.set_as_default,
+        )
+        return await impl.update_prompt(request)
+
+    @router.delete(
+        "/prompts/{prompt_id}",
+        summary="Delete a prompt.",
+        description="Delete a prompt.",
+        responses={
+            200: {"description": "The prompt was successfully deleted."},
+        },
+    )
+    async def delete_prompt(
+        request: Annotated[DeletePromptRequest, Depends(delete_prompt_request)],
+    ) -> None:
+        return await impl.delete_prompt(request)
+
+    @router.put(
+        "/prompts/{prompt_id}/set-default-version",
+        response_model=Prompt,
+        summary="Set prompt version.",
+        description="Set which version of a prompt should be the default in get_prompt (latest).",
+        responses={
+            200: {"description": "The prompt with the specified version now set as default."},
+        },
+    )
+    async def set_default_version(
+        prompt_id: Annotated[str, Path(description="The identifier of the prompt.")],
+        body: Annotated[SetDefaultVersionBodyRequest, Body(...)],
+    ) -> Prompt:
+        request = SetDefaultVersionRequest(prompt_id=prompt_id, version=body.version)
+        return await impl.set_default_version(request)
+
+    return router