From 7b93964a164f0103c25b25d4fc4941f44574d272 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Han?= <seb@redhat.com>
Date: Tue, 2 Dec 2025 15:19:41 +0100
Subject: [PATCH] chore: extract the protocol into its own file
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The protocol leaves in api.py now

Signed-off-by: Sébastien Han <seb@redhat.com>
---
 src/llama_stack_api/batches/__init__.py | 42 ++-----------------
 src/llama_stack_api/batches/api.py      | 56 +++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 39 deletions(-)
 create mode 100644 src/llama_stack_api/batches/api.py

diff --git a/src/llama_stack_api/batches/__init__.py b/src/llama_stack_api/batches/__init__.py
index 33dc62b81..4ad36663a 100644
--- a/src/llama_stack_api/batches/__init__.py
+++ b/src/llama_stack_api/batches/__init__.py
@@ -11,13 +11,14 @@ Pydantic models are defined in llama_stack_api.batches.models.
 The FastAPI router is defined in llama_stack_api.batches.fastapi_routes.
 """
 
-from typing import Protocol, runtime_checkable
-
 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
 
+# Import protocol for re-export
+from llama_stack_api.batches.api import Batches
+
 # Import models for re-export
 from llama_stack_api.batches.models import (
     CancelBatchRequest,
@@ -27,43 +28,6 @@ from llama_stack_api.batches.models import (
     RetrieveBatchRequest,
 )
 
-
-@runtime_checkable
-class Batches(Protocol):
-    """
-    The Batches API enables efficient processing of multiple requests in a single operation,
-    particularly useful for processing large datasets, batch evaluation workflows, and
-    cost-effective inference at scale.
-
-    The API is designed to allow use of openai client libraries for seamless integration.
-
-    This API provides the following extensions:
-     - idempotent batch creation
-
-    Note: This API is currently under active development and may undergo changes.
-    """
-
-    async def create_batch(
-        self,
-        request: CreateBatchRequest,
-    ) -> BatchObject: ...
-
-    async def retrieve_batch(
-        self,
-        request: RetrieveBatchRequest,
-    ) -> BatchObject: ...
-
-    async def cancel_batch(
-        self,
-        request: CancelBatchRequest,
-    ) -> BatchObject: ...
-
-    async def list_batches(
-        self,
-        request: ListBatchesRequest,
-    ) -> ListBatchesResponse: ...
-
-
 __all__ = [
     "Batches",
     "BatchObject",
diff --git a/src/llama_stack_api/batches/api.py b/src/llama_stack_api/batches/api.py
new file mode 100644
index 000000000..2f920fe07
--- /dev/null
+++ b/src/llama_stack_api/batches/api.py
@@ -0,0 +1,56 @@
+# 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 typing import Protocol, runtime_checkable
+
+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
+
+from llama_stack_api.batches.models import (
+    CancelBatchRequest,
+    CreateBatchRequest,
+    ListBatchesRequest,
+    ListBatchesResponse,
+    RetrieveBatchRequest,
+)
+
+
+@runtime_checkable
+class Batches(Protocol):
+    """
+    The Batches API enables efficient processing of multiple requests in a single operation,
+    particularly useful for processing large datasets, batch evaluation workflows, and
+    cost-effective inference at scale.
+
+    The API is designed to allow use of openai client libraries for seamless integration.
+
+    This API provides the following extensions:
+     - idempotent batch creation
+
+    Note: This API is currently under active development and may undergo changes.
+    """
+
+    async def create_batch(
+        self,
+        request: CreateBatchRequest,
+    ) -> BatchObject: ...
+
+    async def retrieve_batch(
+        self,
+        request: RetrieveBatchRequest,
+    ) -> BatchObject: ...
+
+    async def cancel_batch(
+        self,
+        request: CancelBatchRequest,
+    ) -> BatchObject: ...
+
+    async def list_batches(
+        self,
+        request: ListBatchesRequest,
+    ) -> ListBatchesResponse: ...