From 2fe24a6df8d6b650cdfc64843f3e6fd593cceba6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Han?= Date: Thu, 20 Nov 2025 12:41:24 +0100 Subject: [PATCH] chore: move ListBatchesResponse to models.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sébastien Han --- .../core/server/routers/batches.py | 4 ---- src/llama_stack_api/batches/__init__.py | 19 ++++-------------- src/llama_stack_api/batches/models.py | 20 ++++++++++++++++--- 3 files changed, 21 insertions(+), 22 deletions(-) diff --git a/src/llama_stack/core/server/routers/batches.py b/src/llama_stack/core/server/routers/batches.py index fb7f8ebfa..de1605eaf 100644 --- a/src/llama_stack/core/server/routers/batches.py +++ b/src/llama_stack/core/server/routers/batches.py @@ -56,7 +56,6 @@ def create_batches_router(impl_getter: Callable[[Api], Batches]) -> APIRouter: request: Annotated[CreateBatchRequest, Body(...)], svc: Annotated[Batches, Depends(get_batch_service)], ) -> BatchObject: - """Create a new batch.""" return await svc.create_batch( input_file_id=request.input_file_id, endpoint=request.endpoint, @@ -78,7 +77,6 @@ def create_batches_router(impl_getter: Callable[[Api], Batches]) -> APIRouter: batch_id: str, svc: Annotated[Batches, Depends(get_batch_service)], ) -> BatchObject: - """Retrieve information about a specific batch.""" return await svc.retrieve_batch(batch_id) @router.post( @@ -94,7 +92,6 @@ def create_batches_router(impl_getter: Callable[[Api], Batches]) -> APIRouter: batch_id: str, svc: Annotated[Batches, Depends(get_batch_service)], ) -> BatchObject: - """Cancel a batch that is in progress.""" return await svc.cancel_batch(batch_id) @router.get( @@ -111,7 +108,6 @@ def create_batches_router(impl_getter: Callable[[Api], Batches]) -> APIRouter: after: str | None = None, limit: int = 20, ) -> ListBatchesResponse: - """List all batches for the current user.""" return await svc.list_batches(after=after, limit=limit) return router diff --git a/src/llama_stack_api/batches/__init__.py b/src/llama_stack_api/batches/__init__.py index 636dd0c52..2dff546b4 100644 --- a/src/llama_stack_api/batches/__init__.py +++ b/src/llama_stack_api/batches/__init__.py @@ -6,31 +6,20 @@ """Batches API protocol and models. -This module contains the Batches protocol definition and related models. +This module contains the Batches protocol definition. +Pydantic models are defined in llama_stack_api.batches.models. The router implementation is in llama_stack.core.server.routers.batches. """ from typing import Literal, Protocol, runtime_checkable -from pydantic import BaseModel, Field - -from llama_stack_api.schema_utils import json_schema_type - try: from openai.types import Batch as BatchObject except ImportError as e: raise ImportError("OpenAI package is required for batches API. Please install it with: pip install openai") from e - -@json_schema_type -class ListBatchesResponse(BaseModel): - """Response containing a list of batch objects.""" - - object: Literal["list"] = "list" - data: list[BatchObject] = Field(..., description="List of batch objects") - first_id: str | None = Field(default=None, description="ID of the first batch in the list") - last_id: str | None = Field(default=None, description="ID of the last batch in the list") - has_more: bool = Field(default=False, description="Whether there are more batches available") +# Import models for re-export +from llama_stack_api.batches.models import ListBatchesResponse @runtime_checkable diff --git a/src/llama_stack_api/batches/models.py b/src/llama_stack_api/batches/models.py index 22e024be2..fe449280d 100644 --- a/src/llama_stack_api/batches/models.py +++ b/src/llama_stack_api/batches/models.py @@ -14,9 +14,13 @@ from typing import Literal from pydantic import BaseModel, Field -from llama_stack_api.batches import BatchObject, ListBatchesResponse from llama_stack_api.schema_utils import json_schema_type +try: + from openai.types import Batch as BatchObject +except ImportError as e: + raise ImportError("OpenAI package is required for batches API. Please install it with: pip install openai") from e + @json_schema_type class CreateBatchRequest(BaseModel): @@ -33,5 +37,15 @@ class CreateBatchRequest(BaseModel): ) -# Re-export response models for convenience -__all__ = ["CreateBatchRequest", "BatchObject", "ListBatchesResponse"] +@json_schema_type +class ListBatchesResponse(BaseModel): + """Response containing a list of batch objects.""" + + object: Literal["list"] = "list" + data: list[BatchObject] = Field(..., description="List of batch objects") + first_id: str | None = Field(default=None, description="ID of the first batch in the list") + last_id: str | None = Field(default=None, description="ID of the last batch in the list") + has_more: bool = Field(default=False, description="Whether there are more batches available") + + +__all__ = ["CreateBatchRequest", "ListBatchesResponse", "BatchObject"]