from openai import OpenAI
import openai
import time
import json
import os


class OpenAIBatchProcessor:
    def __init__(self, api_key):
        # client = OpenAI(api_key=api_key)
        client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="EMPTY")

        self.client = client

    def process_batch(self, input_file_path, endpoint, completion_window):

        # Upload the input file
        with open(input_file_path, "rb") as file:
            uploaded_file = self.client.files.create(file=file, purpose="batch")

        # Create the batch job
        batch_job = self.client.batches.create(
            input_file_id=uploaded_file.id,
            endpoint=endpoint,
            completion_window=completion_window,
        )

        # Monitor the batch job status
        while batch_job.status not in ["completed", "failed", "cancelled"]:
            time.sleep(3)  # Wait for 3 seconds before checking the status again
            print(
                f"Batch job status: {batch_job.status}...trying again in 3 seconds..."
            )
            batch_job = self.client.batches.retrieve(batch_job.id)

        # Check the batch job status and errors
        if batch_job.status == "failed":
            print(f"Batch job failed with status: {batch_job.status}")
            print(f"Batch job errors: {batch_job.errors}")
            return None

        # If the batch job is completed, process the results
        if batch_job.status == "completed":

            # print result of batch job
            print("batch", batch_job.request_counts)

            result_file_id = batch_job.output_file_id
            # Retrieve the file content from the server
            file_response = self.client.files.content(result_file_id)
            result_content = file_response.read()  # Read the content of the file

            # Save the content to a local file
            result_file_name = "batch_job_chat_results.jsonl"
            with open(result_file_name, "wb") as file:
                file.write(result_content)  # Write the binary content to the file
            # Load data from the saved JSONL file
            results = []
            with open(result_file_name, "r", encoding="utf-8") as file:
                for line in file:
                    json_object = json.loads(
                        line.strip()
                    )  # Parse each line as a JSON object
                    results.append(json_object)

            return results
        else:
            print(f"Batch job failed with status: {batch_job.status}")
            return None


# Initialize the OpenAIBatchProcessor
api_key = os.environ.get("OPENAI_API_KEY")
processor = OpenAIBatchProcessor(api_key)

# Process the batch job
input_file_path = "input.jsonl"
endpoint = "/v1/chat/completions"
completion_window = "24h"

# Process the batch job
results = processor.process_batch(input_file_path, endpoint, completion_window)

# Print the results
print(results)

Batch job status: validating...trying again in 3 seconds...
Batch job status: in_progress...trying again in 3 seconds...
Batch job status: in_progress...trying again in 3 seconds...
batch BatchRequestCounts(completed=2, failed=0, total=2)
[{'id': 'batch_req_bee263f2-0734-44cb-96c4-7fb8c4ad1661', 'custom_id': 'request-1', 'response': {'status_code': 200, 'request_id': '2a93b0ea158c42faa4efb9578ad75064', 'body': {'id': '2a93b0ea158c42faa4efb9578ad75064', 'object': 'chat.completion', 'created': 1721670579, 'model': 'gpt-3.5-turbo-0125', 'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': "  Hello there! 👋 It's my pleasure to assist you. Here are three NBA players:\n\n1. LeBron James - He is a four-time NBA champion and four-time NBA Most Valuable Player (MVP) who has played for the Cleveland Cavaliers, Miami Heat, and Los Angeles Lakers.\n2. Stephen Curry - He is a three-time NBA champion and two-time NBA MVP who has played for the Golden State Warriors. He is known for his incredible shooting ability and has won multiple awards for his skills.\n3. Kevin Durant - He is a two-time NBA champion and two-time NBA MVP who has played for the Oklahoma City Thunder and Golden State Warriors. He is known for his scoring ability and is considered one of the best players in the NBA.\n\nI hope this helps! Let me know if you have any other questions. 😊"}, 'logprobs': None, 'finish_reason': 'FINISH_MATCHED_TOKEN: 2'}], 'usage': {'prompt_tokens': 37, 'completion_tokens': 203, 'total_tokens': 240}, 'system_fingerprint': None}}, 'error': None}, {'id': 'batch_req_db9eff3b-aead-4641-bb7b-8fba9fc2626c', 'custom_id': 'request-2', 'response': {'status_code': 200, 'request_id': '3c4bf3e052184f7a8f3e37dc71c79496', 'body': {'id': '3c4bf3e052184f7a8f3e37dc71c79496', 'object': 'chat.completion', 'created': 1721670581, 'model': 'gpt-3.5-turbo-0125', 'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': "  Hello there! As an assistant, I'm happy to help. Here are three capital cities:\n\n1. Tokyo, Japan\n2. New York City, USA\n3. London, United Kingdom"}, 'logprobs': None, 'finish_reason': 'FINISH_MATCHED_TOKEN: 2'}], 'usage': {'prompt_tokens': 34, 'completion_tokens': 44, 'total_tokens': 78}, 'system_fingerprint': None}}, 'error': None}]