Add support for Batch API (#325)

Author: rajasbansal

src/together/client.py

@@ -23,6 +23,7 @@ class Together:
     fine_tuning: resources.FineTuning
     rerank: resources.Rerank
     audio: resources.Audio
+    batches: resources.Batches
     code_interpreter: CodeInterpreter
 
     # client options
@@ -90,6 +91,7 @@ def __init__(
         self.audio = resources.Audio(self.client)
         self.endpoints = resources.Endpoints(self.client)
         self.code_interpreter = CodeInterpreter(self.client)
+        self.batches = resources.Batches(self.client)
 
 
 class AsyncTogether:
@@ -102,7 +104,7 @@ class AsyncTogether:
     fine_tuning: resources.AsyncFineTuning
     rerank: resources.AsyncRerank
     code_interpreter: CodeInterpreter
-
+    batches: resources.AsyncBatches
     # client options
     client: TogetherClient
 
@@ -166,6 +168,7 @@ def __init__(
         self.fine_tuning = resources.AsyncFineTuning(self.client)
         self.rerank = resources.AsyncRerank(self.client)
         self.code_interpreter = CodeInterpreter(self.client)
+        self.batches = resources.AsyncBatches(self.client)
 
 
 Client = Together

src/together/resources/__init__.py

@@ -8,6 +8,7 @@
 from together.resources.images import AsyncImages, Images
 from together.resources.models import AsyncModels, Models
 from together.resources.rerank import AsyncRerank, Rerank
+from together.resources.batch import Batches, AsyncBatches
 
 
 __all__ = [
@@ -31,4 +32,6 @@
     "Audio",
     "AsyncEndpoints",
     "Endpoints",
+    "Batches",
+    "AsyncBatches",
 ]

src/together/resources/batch.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from typing import List
+
+from together.abstract import api_requestor
+from together.together_response import TogetherResponse
+from together.types import (
+    TogetherClient,
+    TogetherRequest,
+    BatchJob,
+)
+
+
+class Batches:
+    def __init__(self, client: TogetherClient) -> None:
+        self._client = client
+
+    def create_batch(self, file_id: str, endpoint: str) -> BatchJob:
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = {
+            "input_file_id": file_id,
+            "endpoint": endpoint,
+            "completion_window": "24h",
+        }
+
+        response, _, _ = requestor.request(
+            options=TogetherRequest(
+                method="POST",
+                url=f"batches",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, TogetherResponse)
+        response_body = response.data.get("job", {})
+        return BatchJob(**response_body)
+
+    def get_batch(self, batch_job_id: str) -> BatchJob:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=TogetherRequest(
+                method="GET",
+                url=f"batches/{batch_job_id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, TogetherResponse)
+        return BatchJob(**response.data)
+
+    def list_batches(self) -> List[BatchJob]:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=TogetherRequest(
+                method="GET",
+                url="batches",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, TogetherResponse)
+        jobs = response.data or []
+        return [BatchJob(**job) for job in jobs]
+
+
+class AsyncBatches:
+    def __init__(self, client: TogetherClient) -> None:
+        self._client = client
+
+    async def create_batch(self, file_id: str, endpoint: str) -> BatchJob:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = {
+            "input_file_id": file_id,
+            "endpoint": endpoint,
+            "completion_window": "24h",
+        }
+
+        response, _, _ = await requestor.arequest(
+            options=TogetherRequest(
+                method="POST",
+                url=f"batches",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, TogetherResponse)
+        response_body = response.data.get("job", {})
+        return BatchJob(**response_body)
+
+    async def get_batch(self, batch_job_id: str) -> BatchJob:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=TogetherRequest(
+                method="GET",
+                url=f"batches/{batch_job_id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, TogetherResponse)
+        return BatchJob(**response.data)
+
+    async def list_batches(self) -> List[BatchJob]:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=TogetherRequest(
+                method="GET",
+                url="batches",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, TogetherResponse)
+        jobs = response.data or []
+        return [BatchJob(**job) for job in jobs]

src/together/resources/files.py

@@ -32,7 +32,7 @@ def upload(
     ) -> FileResponse:
         upload_manager = UploadManager(self._client)
 
-        if check:
+        if check and purpose == FilePurpose.FineTune:
             report_dict = check_file(file)
             if not report_dict["is_check_passed"]:
                 raise FileTypeError(

src/together/types/__init__.py

@@ -52,6 +52,7 @@
 from together.types.images import ImageRequest, ImageResponse
 from together.types.models import ModelObject
 from together.types.rerank import RerankRequest, RerankResponse
+from together.types.batch import BatchJob, BatchJobStatus, BatchEndpoint
 
 
 __all__ = [
@@ -104,4 +105,7 @@
     "DedicatedEndpoint",
     "ListEndpoint",
     "Autoscaling",
+    "BatchJob",
+    "BatchJobStatus",
+    "BatchEndpoint",
 ]

src/together/types/batch.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Optional
+from datetime import datetime
+
+from pydantic import Field
+
+from together.types.abstract import BaseModel
+
+
+class BatchJobStatus(str, Enum):
+    """
+    The status of a batch job
+    """
+
+    VALIDATING = "VALIDATING"
+    IN_PROGRESS = "IN_PROGRESS"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+    EXPIRED = "EXPIRED"
+    CANCELLED = "CANCELLED"
+
+
+class BatchEndpoint(str, Enum):
+    """
+    The endpoint of a batch job
+    """
+
+    COMPLETIONS = "/v1/completions"
+    CHAT_COMPLETIONS = "/v1/chat/completions"
+    # More endpoints can be added here as needed
+
+
+class BatchJob(BaseModel):
+    """
+    A batch job object
+    """
+
+    id: str
+    user_id: str
+    input_file_id: str
+    file_size_bytes: int
+    status: BatchJobStatus
+    job_deadline: datetime
+    created_at: datetime
+    endpoint: str
+    progress: float = 0.0
+    model_id: Optional[str] = None
+    output_file_id: Optional[str] = None
+    error_file_id: Optional[str] = None
+    error: Optional[str] = None
+    completed_at: Optional[datetime] = None

src/together/types/files.py

@@ -13,6 +13,7 @@
 
 class FilePurpose(str, Enum):
     FineTune = "fine-tune"
+    BatchAPI = "batch-api"
 
 
 class FileType(str, Enum):