import os
import json
from openai import OpenAI

# 初始化OpenAI客户端
client = OpenAI(
    api_key="sk-xxx",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# 获取用户输入的文件夹路径
folder_path = input("请输入包含批量文件的文件夹路径：")
if not os.path.isdir(folder_path):
    print("错误：输入的路径不存在或不是文件夹")
    exit()

# 获取文件夹下所有文件列表
file_list = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]

# 创建并保存batch任务信息
with open("batch_results.jsonl", "w", encoding="utf-8") as jsonl_file:
    for filename in file_list:
        # 构建OSS文件标识符（假设OSS路径结构与本地一致）
        input_file_id = f"oss:me-east-1:acm-mm/test/{filename}"
        
        try:
            # 创建batch任务
            batch = client.batches.create(
                input_file_id=input_file_id,
                endpoint="/v1/chat/completions",
                completion_window="1h"
            )
            
            # 提取需要的信息
            batch_info = {
                "id": batch.id,
                "input_file_id": batch.input_file_id
            }
            
            # 写入JSONL文件
            jsonl_file.write(json.dumps(batch_info, ensure_ascii=False) + "\n")
            print(f"文件 {filename} 已成功创建任务，ID: {batch.id}")
            
        except Exception as e:
            print(f"为文件 {filename} 创建任务失败，错误信息: {str(e)}")