243 lines
10 KiB
Python
243 lines
10 KiB
Python
"""
|
|
OpenAI-specific batch processing implementation.
|
|
|
|
This module contains the OpenAI batch processing provider class.
|
|
"""
|
|
|
|
from typing import Any, Optional, Union
|
|
import io
|
|
import logging
|
|
from .base import BatchProvider
|
|
from ..models import BatchJobInfo
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class OpenAIProvider(BatchProvider):
|
|
"""OpenAI batch processing provider"""
|
|
|
|
def submit_batch(
|
|
self,
|
|
file_path_or_buffer: Union[str, io.BytesIO],
|
|
metadata: Optional[dict[str, Any]] = None,
|
|
**kwargs,
|
|
) -> str:
|
|
"""Submit OpenAI batch job"""
|
|
try:
|
|
from openai import OpenAI
|
|
|
|
client = OpenAI()
|
|
|
|
if metadata is None:
|
|
metadata = {"description": "Instructor batch job"}
|
|
|
|
logger.debug(f"Submitting batch job with metadata: {metadata}")
|
|
|
|
if isinstance(file_path_or_buffer, str):
|
|
logger.debug(f"Creating batch file from path: {file_path_or_buffer}")
|
|
with open(file_path_or_buffer, "rb") as f:
|
|
batch_file = client.files.create(file=f, purpose="batch")
|
|
elif isinstance(file_path_or_buffer, io.BytesIO):
|
|
logger.debug("Creating batch file from BytesIO buffer")
|
|
file_path_or_buffer.seek(0)
|
|
batch_file = client.files.create(
|
|
file=file_path_or_buffer, purpose="batch"
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported file_path_or_buffer type: {type(file_path_or_buffer)}"
|
|
)
|
|
|
|
batch_job = client.batches.create(
|
|
input_file_id=batch_file.id,
|
|
endpoint="/v1/chat/completions",
|
|
completion_window=kwargs.get("completion_window", "24h"),
|
|
metadata=metadata,
|
|
)
|
|
logger.info(f"Successfully submitted batch job: {batch_job.id}")
|
|
return batch_job.id
|
|
except (ValueError, TypeError) as e:
|
|
# Re-raise validation errors as-is
|
|
logger.error(f"Validation error in OpenAI batch submission: {e}")
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Failed to submit OpenAI batch: {e}")
|
|
raise RuntimeError(f"Failed to submit OpenAI batch: {e}") from e
|
|
|
|
def get_status(self, batch_id: str) -> dict[str, Any]:
|
|
"""Get OpenAI batch status"""
|
|
try:
|
|
from openai import OpenAI
|
|
|
|
client = OpenAI()
|
|
batch = client.batches.retrieve(batch_id)
|
|
return {
|
|
"id": batch.id,
|
|
"status": batch.status,
|
|
"created_at": batch.created_at,
|
|
"request_counts": {
|
|
"total": getattr(batch.request_counts, "total", 0),
|
|
"completed": getattr(batch.request_counts, "completed", 0),
|
|
"failed": getattr(batch.request_counts, "failed", 0),
|
|
},
|
|
}
|
|
except Exception as e:
|
|
raise Exception(f"Failed to get OpenAI batch status: {e}") from e
|
|
|
|
def retrieve_results(self, batch_id: str) -> str:
|
|
"""Retrieve OpenAI batch results"""
|
|
try:
|
|
from openai import OpenAI
|
|
import time
|
|
|
|
client = OpenAI()
|
|
batch = client.batches.retrieve(batch_id)
|
|
|
|
if batch.status != "completed":
|
|
raise Exception(f"Batch not completed, status: {batch.status}")
|
|
|
|
# Check if all requests failed
|
|
request_counts = getattr(batch, "request_counts", None)
|
|
if request_counts:
|
|
completed = getattr(request_counts, "completed", 0)
|
|
failed = getattr(request_counts, "failed", 0)
|
|
total = getattr(request_counts, "total", 0)
|
|
|
|
if failed > 0 and completed == 0:
|
|
raise RuntimeError(
|
|
f"All {total} batch requests failed. No output file will be available. "
|
|
)
|
|
|
|
if not batch.output_file_id:
|
|
# Sometimes output file isn't immediately available, wait longer and retry more
|
|
max_retries = 10
|
|
for attempt in range(max_retries):
|
|
wait_time = min(
|
|
5 + attempt, 15
|
|
) # Progressive backoff: 5s, 6s, 7s... up to 15s
|
|
print(
|
|
f"Output file not ready, waiting {wait_time}s (attempt {attempt + 1}/{max_retries})..."
|
|
)
|
|
time.sleep(wait_time)
|
|
batch = client.batches.retrieve(batch_id)
|
|
if batch.output_file_id:
|
|
print(f"Output file now available: {batch.output_file_id}")
|
|
break
|
|
# Check if batch failed during our wait
|
|
if batch.status != "completed":
|
|
raise Exception(
|
|
f"Batch status changed to {batch.status} while waiting for output file"
|
|
)
|
|
if attempt == max_retries - 1:
|
|
# Final attempt - provide detailed error info
|
|
raise RuntimeError(
|
|
f"No output file available after {max_retries} retries over {sum(range(5, 5 + max_retries))} seconds. "
|
|
f"Batch status: {batch.status}, Request counts: {getattr(batch, 'request_counts', 'unknown')}. "
|
|
)
|
|
|
|
if batch.output_file_id is None:
|
|
raise RuntimeError("Batch has no output file ID available")
|
|
file_response = client.files.content(batch.output_file_id)
|
|
return file_response.text
|
|
except Exception as e:
|
|
raise Exception(f"Failed to retrieve OpenAI results: {e}") from e
|
|
|
|
def download_results(self, batch_id: str, file_path: str) -> None:
|
|
"""Download OpenAI batch results to a file"""
|
|
try:
|
|
from openai import OpenAI
|
|
import time
|
|
|
|
client = OpenAI()
|
|
batch = client.batches.retrieve(batch_id)
|
|
|
|
if batch.status != "completed":
|
|
raise Exception(f"Batch not completed, status: {batch.status}")
|
|
|
|
# Check if all requests failed
|
|
request_counts = getattr(batch, "request_counts", None)
|
|
if request_counts:
|
|
completed = getattr(request_counts, "completed", 0)
|
|
failed = getattr(request_counts, "failed", 0)
|
|
total = getattr(request_counts, "total", 0)
|
|
|
|
if failed > 0 and completed == 0:
|
|
raise RuntimeError(
|
|
f"All {total} batch requests failed. No output file will be available."
|
|
)
|
|
|
|
if not batch.output_file_id:
|
|
# Sometimes output file isn't immediately available, wait longer and retry more
|
|
max_retries = 10
|
|
for attempt in range(max_retries):
|
|
wait_time = min(
|
|
5 + attempt, 15
|
|
) # Progressive backoff: 5s, 6s, 7s... up to 15s
|
|
print(
|
|
f"Output file not ready, waiting {wait_time}s (attempt {attempt + 1}/{max_retries})..."
|
|
)
|
|
time.sleep(wait_time)
|
|
batch = client.batches.retrieve(batch_id)
|
|
if batch.output_file_id:
|
|
print(f"Output file now available: {batch.output_file_id}")
|
|
break
|
|
# Check if batch failed during our wait
|
|
if batch.status != "completed":
|
|
raise Exception(
|
|
f"Batch status changed to {batch.status} while waiting for output file"
|
|
)
|
|
if attempt == max_retries - 1:
|
|
# Final attempt - provide detailed error info
|
|
raise Exception(
|
|
f"No output file available after {max_retries} retries over {sum(range(5, 5 + max_retries))} seconds. "
|
|
f"Batch status: {batch.status}, Request counts: {getattr(batch, 'request_counts', 'unknown')}."
|
|
)
|
|
|
|
if batch.output_file_id is None:
|
|
raise RuntimeError("Batch has no output file ID available")
|
|
file_response = client.files.content(batch.output_file_id)
|
|
with open(file_path, "w") as f:
|
|
f.write(file_response.text)
|
|
except Exception as e:
|
|
raise Exception(f"Failed to download OpenAI results: {e}") from e
|
|
|
|
def cancel_batch(self, batch_id: str) -> dict[str, Any]:
|
|
"""Cancel OpenAI batch job"""
|
|
try:
|
|
from openai import OpenAI
|
|
|
|
client = OpenAI()
|
|
batch = client.batches.cancel(batch_id)
|
|
return batch.model_dump()
|
|
except Exception as e:
|
|
raise Exception(f"Failed to cancel OpenAI batch: {e}") from e
|
|
|
|
def delete_batch(self, batch_id: str) -> dict[str, Any]:
|
|
"""Delete OpenAI batch job"""
|
|
try:
|
|
from openai import OpenAI
|
|
|
|
client = OpenAI()
|
|
# OpenAI doesn't have a delete endpoint, so we'll return the batch info
|
|
batch = client.batches.retrieve(batch_id)
|
|
return {
|
|
"id": batch.id,
|
|
"status": batch.status,
|
|
"message": "OpenAI does not support batch deletion",
|
|
}
|
|
except Exception as e:
|
|
raise Exception(f"Failed to delete OpenAI batch: {e}") from e
|
|
|
|
def list_batches(self, limit: int = 10) -> list[BatchJobInfo]:
|
|
"""List OpenAI batch jobs"""
|
|
try:
|
|
from openai import OpenAI
|
|
|
|
client = OpenAI()
|
|
batches = client.batches.list(limit=limit)
|
|
return [
|
|
BatchJobInfo.from_openai(batch.model_dump()) for batch in batches.data
|
|
]
|
|
except Exception as e:
|
|
raise Exception(f"Failed to list OpenAI batches: {e}") from e
|