diff --git a/src/together/client.py b/src/together/client.py index 08d829d1..a1784529 100644 --- a/src/together/client.py +++ b/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 diff --git a/src/together/resources/__init__.py b/src/together/resources/__init__.py index f07aeb00..a6b1935f 100644 --- a/src/together/resources/__init__.py +++ b/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", ] diff --git a/src/together/resources/batch.py b/src/together/resources/batch.py new file mode 100644 index 00000000..9602af2b --- /dev/null +++ b/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] diff --git a/src/together/resources/files.py b/src/together/resources/files.py index 14500b24..3f637c8a 100644 --- a/src/together/resources/files.py +++ b/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( diff --git a/src/together/types/__init__.py b/src/together/types/__init__.py index fddb3636..e218a02e 100644 --- a/src/together/types/__init__.py +++ b/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", ] diff --git a/src/together/types/batch.py b/src/together/types/batch.py new file mode 100644 index 00000000..3bd5ead5 --- /dev/null +++ b/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 diff --git a/src/together/types/files.py b/src/together/types/files.py index 24a6cd99..59acd8c6 100644 --- a/src/together/types/files.py +++ b/src/together/types/files.py @@ -13,6 +13,7 @@ class FilePurpose(str, Enum): FineTune = "fine-tune" + BatchAPI = "batch-api" class FileType(str, Enum):