Skip to content

Add support for Batch API #325

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 11, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion src/together/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ class Together:
fine_tuning: resources.FineTuning
rerank: resources.Rerank
audio: resources.Audio
batches: resources.Batches
code_interpreter: CodeInterpreter

# client options
Expand Down Expand Up @@ -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:
Expand All @@ -102,7 +104,7 @@ class AsyncTogether:
fine_tuning: resources.AsyncFineTuning
rerank: resources.AsyncRerank
code_interpreter: CodeInterpreter

batches: resources.AsyncBatches
# client options
client: TogetherClient

Expand Down Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions src/together/resources/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__ = [
Expand All @@ -31,4 +32,6 @@
"Audio",
"AsyncEndpoints",
"Endpoints",
"Batches",
"AsyncBatches",
]
136 changes: 136 additions & 0 deletions src/together/resources/batch.py
Original file line number Diff line number Diff line change
@@ -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]
2 changes: 1 addition & 1 deletion src/together/resources/files.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
4 changes: 4 additions & 0 deletions src/together/types/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__ = [
Expand Down Expand Up @@ -104,4 +105,7 @@
"DedicatedEndpoint",
"ListEndpoint",
"Autoscaling",
"BatchJob",
"BatchJobStatus",
"BatchEndpoint",
]
53 changes: 53 additions & 0 deletions src/together/types/batch.py
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions src/together/types/files.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

class FilePurpose(str, Enum):
FineTune = "fine-tune"
BatchAPI = "batch-api"


class FileType(str, Enum):
Expand Down