Привет короче есть у меня такой код и проблема что батчи могут обрабатываться до 24ч и проблема такая что мне нужно заходить на другой день запускать программу и она должна запоминать какие батчи она ждет и все данных сохранять для этого и после запуска снова ожидать батчи, и когда батч закончится корректно его обработать как вот в этом коде у меня. короче сделай это import openai

import base64

import requests

import os

import csv

import json

import re

from configparser import ConfigParser

import time

import logging

from PIL import Image, ImageOps

import shutil

from tqdm import tqdm

# Настройка логирования

logging.basicConfig(

    level=logging.INFO,

    format='%(asctime)s %(levelname)s:%(message)s',

    handlers=[

        logging.FileHandler("batch_processing.log"),

        logging.StreamHandler()

    ]

)

# Настройка OpenAI API

def load_api_key(config_path='config.ini'):

    config = ConfigParser()

    config.read(config_path)

    return config.get('openai', 'api_key')

# Функция для кодирования изображения в base64

def encode_image(image_path):

    with open(image_path, "rb") as image_file:

        return base64.b64encode(image_file.read()).decode('utf-8')

# Функция для создания задач для Batch API с уникальными custom_id

def create_batch_tasks(image_files, directory_path, prompt_template, starting_index=0):

    tasks = []

    for index, filename in enumerate(image_files):

        image_path = os.path.join(directory_path, filename)

        base64_image = encode_image(image_path)

       

        task = {

            "custom_id": f"task-{starting_index + index}",  # Уникальный идентификатор

            "method": "POST",

            "url": "/v1/chat/completions",

            "body": {

                "model": "gpt-4o",

                "messages": [

                    {

                        "role": "user",

                        "content": [

                            {"type": "text", "text": prompt_template},

                            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}", "detail": "high"}}

                        ]

                    }

                ],

                "max_tokens": 300,

                "response_format": {

                    "type": "json_object"

                }

            }

        }

        tasks.append(task)

    return tasks

# Функция для создания JSONL файла

def write_jsonl_file(tasks, file_path):

    with open(file_path, 'w', encoding='utf-8') as f:

        for task in tasks:

            f.write(json.dumps(task, ensure_ascii=False) + '\n')

# Функция для загрузки файла на OpenAI

def upload_file(file_path, api_key):

    headers = {

        "Authorization": f"Bearer {api_key}"

    }

    with open(file_path, "rb") as f:

        response = requests.post(

            "https://api.openai.com/v1/files",

            headers=headers,

            files={

                'file': (os.path.basename(file_path), f, 'application/jsonl')

            },

            data={

                'purpose': 'batch'

            }

        )

    response.raise_for_status()

    return response.json()['id']

# Функция для создания Batch Job

def create_batch_job(file_id, api_key):

    headers = {

        "Authorization": f"Bearer {api_key}",

        "Content-Type": "application/json"

    }

    payload = {

        "input_file_id": file_id,

        "endpoint": "/v1/chat/completions",

        "completion_window": "24h"  # Максимальное время выполнения

    }

    response = requests.post(

        "https://api.openai.com/v1/batches",

        headers=headers,

        json=payload

    )

    response.raise_for_status()

    return response.json()['id']

# Функция для проверки статуса Batch Job

def check_batch_job_status(job_id, api_key):

    headers = {

        "Authorization": f"Bearer {api_key}"

    }

    response = requests.get(

        f"https://api.openai.com/v1/batches/{job_id}",

        headers=headers

    )

    response.raise_for_status()

    return response.json()

# Функция для загрузки результатов Batch Job

def download_batch_results(result_file_id, api_key):

    headers = {

        "Authorization": f"Bearer {api_key}"

    }

    response = requests.get(

        f"https://api.openai.com/v1/files/{result_file_id}/content",

        headers=headers

    )

    response.raise_for_status()

    return response.content

# Функция для парсинга результатов

def parse_results(result_content):

    results = []

    for line in result_content.decode('utf-8').splitlines():

        result = json.loads(line)

        # Извлекаем содержание и использование токенов, если они есть

        response_content = result.get('response', {}).get('body', {})

        choices = response_content.get('choices', [])

        usage = response_content.get('usage', {})

       

        if choices:

            content = choices[0].get('message', {}).get('content', '')

        else:

            content = ''

       

        result_dict = {

            'custom_id': result.get('custom_id', ''),

            'content': content,

            'prompt_tokens': usage.get('prompt_tokens', 0),

            'completion_tokens': usage.get('completion_tokens', 0),

            'total_tokens': usage.get('total_tokens', 0)

        }

        results.append(result_dict)

    return results

# Функция для записи результатов в CSV с форматированием имен и учётом токенов

def write_results_to_csv(results, csv_path, image_files, totals):

    # Создаем словарь для быстрого доступа по custom_id

    results_dict = {res['custom_id']: res for res in results}

    # Открываем CSV файл

    with open(csv_path, 'w', newline='', encoding='utf-8') as csvfile:

        writer = csv.writer(csvfile)

        # Записываем заголовок

        writer.writerow(['Filename', 'Фамилия', 'Имя', 'Отчество', 'Дата рождения', 'Дата смерти', 'Статус', 'Входящие токены', 'Исходящие токены', 'Всего токенов', 'Стоимость ($)'])

        for index, filename in enumerate(image_files):

            task_id = f"task-{index}"

            res = results_dict.get(task_id, None)

            if res and 'content' in res:

                try:

                    content = res['content']

                    # Удаляем тройные обратные кавычки, если они есть

                    content = content.strip('```json').strip('```').strip()

                    parsed = json.loads(content)

                    # Проверяем, является ли parsed списком и содержит ли он элементы

                    if isinstance(parsed, list) and len(parsed) > 0:

                        parsed_data = parsed  # Список, каждый элемент - отдельная могила

                    elif isinstance(parsed, dict):

                        parsed_data = [parsed]  # Преобразуем в список для униформности

                    else:

                        parsed_data = []  # В противном случае, пустой список

                    for data in parsed_data:

                        # Извлекаем данные из parsed_data с форматированием имен

                        last_name = format_name(data.get('фамилия', '-'))

                        first_name = format_name(data.get('имя', '-'))

                        middle_name = format_name(data.get('отчество', '-'))

                        birth_date = data.get('дата рождения', '-')

                        death_date = data.get('дата смерти', '-')

                        status = data.get('статус', 'Обработка')

                        # Извлечение токенов

                        prompt_tokens = res.get('prompt_tokens', 0)

                        completion_tokens = res.get('completion_tokens', 0)

                        total_tokens = res.get('total_tokens', 0)

                        # Расчет стоимости

                        cost = calculate_cost(prompt_tokens, completion_tokens)

                        # Обновление общих сумм

                        totals['prompt_tokens'] += prompt_tokens

                        totals['completion_tokens'] += completion_tokens

                        totals['total_tokens'] += total_tokens

                        totals['cost'] += cost

                        # Запись строки в CSV

                        writer.writerow([

                            filename,

                            last_name,

                            first_name,

                            middle_name,

                            birth_date,

                            death_date,

                            status,

                            prompt_tokens,

                            completion_tokens,

                            total_tokens,

                            f"{cost:.6f}"

                        ])

                        # Вывод информации по изображению

                        print(f"Файл: {filename}")

                        print(f"Токены: Входящие = {prompt_tokens}, Исходящие = {completion_tokens}, Всего = {total_tokens}")

                        print(f"Стоимость: ${cost:.6f}\n")

                except Exception as e:

                    logging.error(f"Ошибка парсинга результата для {filename}: {e}")

                    # Запись пустых значений с статусом 'Обработка'

                    writer.writerow([filename, '-', '-', '-', '-', '-', 'Обработка', 0, 0, 0, "0.000000"])

                    print(f"Файл: {filename}")

                    print(f"Токены: Входящие = 0, Исходящие = 0, Всего = 0")

                    print(f"Стоимость: $0.000000\n")

            else:

                # Если результат отсутствует или возникла ошибка

                writer.writerow([filename, '-', '-', '-', '-', '-', 'Обработка', 0, 0, 0, "0.000000"])

                print(f"Файл: {filename}")

                print(f"Токены: Входящие = 0, Исходящие = 0, Всего = 0")

                print(f"Стоимость: $0.000000\n")

# Функция для форматирования имен: первая буква заглавная, остальные строчные

def format_name(name):

    if name == '-' or not name:

        return '-'

    return name.capitalize()

# Функция для сжатия изображений и удаления метаданных

def compress_and_remove_metadata(input_path, output_path, quality=85):

    try:

        with Image.open(input_path) as img:

            # Корректировка ориентации изображения согласно EXIF

            img = ImageOps.exif_transpose(img)

            # Сжатие изображения и удаление всех метаданных

            img.save(output_path, "JPEG", quality=quality, optimize=True)

    except Exception as e:

        logging.error(f"Ошибка при обработке изображения {input_path}: {e}")

# Функция для загрузки промпта из файла

def load_prompt_from_file(prompt_file_path):

    try:

        with open(prompt_file_path, 'r', encoding='utf-8') as file:

            return file.read()

    except Exception as e:

        logging.error(f"Ошибка загрузки промпта из файла: {e}")

        return None

# Функция для разделения файлов на батчи с ограничением по размеру

def get_batches(file_paths, max_size_mb=100):

    batches = []

    current_batch = []

    current_size = 0

    for file in file_paths:

        size_mb = os.path.getsize(file) / (1024 * 1024)

        if current_size + size_mb > max_size_mb:

            if current_batch:

                batches.append(current_batch)

            current_batch = [file]

            current_size = size_mb

        else:

            current_batch.append(file)

            current_size += size_mb

    if current_batch:

        batches.append(current_batch)

    return batches

# Функция для расчета стоимости на основе токенов

def calculate_cost(prompt_tokens, completion_tokens):

    # Установите цены в зависимости от используемой модели

    # Пример для модели gpt-4o (цены могут изменяться, уточни на сайте OpenAI)

    PROMPT_COST_PER_1000 = 0.00125  # Цена за 1000 входящих токенов

    COMPLETION_COST_PER_1000 = 0.005  # Цена за 1000 исходящих токенов

    cost = (prompt_tokens / 1000) * PROMPT_COST_PER_1000 + (completion_tokens / 1000) * COMPLETION_COST_PER_1000

    return cost

# Основная функция

def main():

    # Загрузка API ключа

    try:

        api_key = load_api_key()

    except Exception as e:

        logging.error(f"Ошибка загрузки API ключа: {e}")

        return

    # Путь к папке с оригинальными изображениями

    original_dir = os.path.join(os.getcwd(), 'photos2')

    # Путь к папке для подготовленных изображений

    prepared_dir = os.path.join(os.getcwd(), 'photos_prepared')

    os.makedirs(prepared_dir, exist_ok=True)

    # Проверка существования папки с изображениями

    if not os.path.exists(original_dir):

        logging.error(f"Папка с изображениями '{original_dir}' не существует.")

        return

    # Список файлов оригинальных изображений

    image_files = [f for f in os.listdir(original_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]

    if not image_files:

        logging.info("Нет файлов для обработки в папке 'photos2'.")

        return

    # Загрузка и сжатие изображений, удаление метаданных

    logging.info("Сжатие изображений и удаление метаданных...")

    prepared_files = []

    for filename in tqdm(image_files, desc="Обработка изображений", unit="изображение"):

        original_path = os.path.join(original_dir, filename)

        prepared_path = os.path.join(prepared_dir, f"prepared_{filename}")

        compress_and_remove_metadata(original_path, prepared_path)

        prepared_files.append(prepared_path)

    # Загрузка промпта из файла

    prompt_file_path = 'prompt.txt'

    prompt_template = load_prompt_from_file(prompt_file_path)

    if not prompt_template:

        logging.error("Промпт не был загружен. Завершение работы.")

        return

    # Разделение подготовленных файлов на батчи

    logging.info("Разделение файлов на батчи...")

    batches = get_batches(prepared_files, max_size_mb=100)

    logging.info(f"Всего батчей: {len(batches)}")

    # Переменная для хранения результатов всех батчей

    all_results = []

    # Инициализация глобального индекса

    global_index = 0

    # Инициализация общих сумм

    totals = {

        'prompt_tokens': 0,

        'completion_tokens': 0,

        'total_tokens': 0,

        'cost': 0.0

    }

    # Создание директории для batch_tasks, если её нет

    os.makedirs('data', exist_ok=True)

    # Обработка каждого батча по очереди

    for batch_num, batch_files in enumerate(batches, start=1):

        logging.info(f"Обработка батча {batch_num}/{len(batches)} с {len(batch_files)} файлами...")

        # Создание задач для текущего батча с использованием глобального индекса

        batch_filenames = [os.path.basename(f) for f in batch_files]

        tasks = create_batch_tasks(batch_filenames, prepared_dir, prompt_template, starting_index=global_index)

        # Обновление глобального индекса

        global_index += len(batch_files)

        # Запись задач в JSONL файл

        tasks_file_path = f"data/batch_tasks_{batch_num}.jsonl"

        os.makedirs(os.path.dirname(tasks_file_path), exist_ok=True)

        logging.info(f"Запись задач в файл {tasks_file_path}...")

        write_jsonl_file(tasks, tasks_file_path)

        # Загрузка файла на OpenAI

        logging.info("Загрузка задач на OpenAI...")

        try:

            batch_file_id = upload_file(tasks_file_path, api_key)

            logging.info(f"Файл батча {batch_num} загружен. ID файла: {batch_file_id}")

        except Exception as e:

            logging.error(f"Ошибка загрузки файла батча {batch_num} на OpenAI: {e}")

            continue  # Переходим к следующему батчу

        # Создание Batch Job

        logging.info("Создание Batch Job...")

        try:

            batch_job_id = create_batch_job(batch_file_id, api_key)

            logging.info(f"Batch Job для батча {batch_num} создан. ID задания: {batch_job_id}")

        except Exception as e:

            logging.error(f"Ошибка создания Batch Job для батча {batch_num}: {e}")

            continue  # Переходим к следующему батчу

        # Отслеживание состояния задания

        logging.info("Отслеживание состояния Batch Job...")

        while True:

            try:

                status_info = check_batch_job_status(batch_job_id, api_key)

                status = status_info.get('status')

                logging.info(f"Текущий статус батча {batch_num}: {status}")

                if status in ['completed', 'failed']:

                    break

                time.sleep(5)  # Ждем 60 секунд перед следующей проверкой

            except Exception as e:

                logging.error(f"Ошибка проверки статуса Batch Job для батча {batch_num}: {e}")

                break  # Переходим к следующему батчу

        if status != 'completed':

            logging.error(f"Batch Job батча {batch_num} завершилось с ошибкой.")

            logging.error(status_info)

            continue  # Переходим к следующему батчу

        logging.info(f"Batch Job батча {batch_num} успешно завершено.")

        # Получение результатов

        result_file_id = status_info.get('output_file_id')

        if not result_file_id:

            logging.error(f"Не найден 'output_file_id' в статусе Batch Job для батча {batch_num}.")

            continue  # Переходим к следующему батчу

        logging.info(f"Загрузка результатов из файла ID: {result_file_id} для батча {batch_num}...")

        try:

            result_content = download_batch_results(result_file_id, api_key)

        except Exception as e:

            logging.error(f"Ошибка загрузки результатов Batch Job для батча {batch_num}: {e}")

            continue  # Переходим к следующему батчу

        # Парсинг результатов

        logging.info("Парсинг результатов...")

        try:

            results = parse_results(result_content)

            all_results.extend(results)

        except Exception as e:

            logging.error(f"Ошибка парсинга результатов батча {batch_num}: {e}")

            continue  # Переходим к следующему батчу

    if not all_results:

        logging.error("Нет результатов для записи в CSV.")

        return

    # Запись всех результатов в единый CSV файл

    csv_output_path = "results.csv"

    logging.info(f"Запись всех результатов в CSV файл: {csv_output_path}...")

    write_results_to_csv(all_results, csv_output_path, image_files, totals)

    # Вывод общих сумм после обработки всех изображений

    print("Общие результаты:")

    print(f"Всего входящих токенов: {totals['prompt_tokens']}")

    print(f"Всего исходящих токенов: {totals['completion_tokens']}")

    print(f"Всего токенов: {totals['total_tokens']}")

    print(f"Общая стоимость: ${totals['cost']:.6f}")

    logging.info("Обработка завершена успешно.")

# Функция для обработки существующего батча

def process_existing_batch(batch_id, api_key, csv_output_path, totals):

    logging.info(f"Обработка существующего Batch ID: {batch_id}")

    # Получение статуса Batch Job

    logging.info("Получение статуса Batch Job...")

    try:

        status_info = check_batch_job_status(batch_id, api_key)

        status = status_info.get('status')

        logging.info(f"Статус Batch Job: {status}")

        if status != 'completed':

            logging.error(f"Batch Job {batch_id} не завершен. Текущий статус: {status}")

            return

    except Exception as e:

        logging.error(f"Ошибка получения статуса Batch Job {batch_id}: {e}")

        return

    # Получение результатов

    result_file_id = status_info.get('output_file_id')

    if not result_file_id:

        logging.error(f"Не найден 'output_file_id' в статусе Batch Job для батча {batch_id}.")

        return

    logging.info(f"Загрузка результатов из файла ID: {result_file_id}...")

    try:

        result_content = download_batch_results(result_file_id, api_key)

    except Exception as e:

        logging.error(f"Ошибка загрузки результатов Batch Job {batch_id}: {e}")

        return

    # Парсинг результатов

    logging.info("Парсинг результатов...")

    try:

        results = parse_results(result_content)

    except Exception as e:

        logging.error(f"Ошибка парсинга результатов Batch Job {batch_id}: {e}")

        return

    # Чтение списка файлов, соответствующих батчу

    # Предполагается, что имена файлов были сохранены или могут быть восстановлены

    # Если нет, нужно дополнительно сохранить соответствие между custom_id и файлами при создании батчей

    # Здесь будет просто пример без привязки к файлам

    image_files = [res['custom_id'] for res in results]

    # Запись результатов в CSV

    logging.info(f"Запись результатов в CSV файл: {csv_output_path}...")

    write_results_to_csv(results, csv_output_path, image_files, totals)

if __name__ == "__main__":

    main()