import os
import json
from openai import OpenAI

# 初始化OpenAI客户端（实际对接阿里云百炼）
client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

def process_batches(input_file):
    output_file = "batch_status_output.jsonl"
    
    with open(input_file, 'r') as infile, open(output_file, 'w') as outfile:
        for line in infile:
            try:
                # 解析JSONL行数据
                entry = json.loads(line.strip())
                batch_id = entry['id']
                
                # 查询Batch详细信息
                batch = client.batches.retrieve(batch_id)
                
                # 构建结果记录
                result = {
                    "status": batch.status,
                    "input_file_id": batch.input_file_id,
                    "output_file_id": batch.output_file_id
                }
                
                # 写入结果到新JSONL文件
                outfile.write(json.dumps(result) + '\n')
                print(f"Processed batch: {batch_id}")
                
            except KeyError:
                print(f"Invalid entry format: {line.strip()}")
            except Exception as e:
                print(f"Error processing batch {batch_id}: {str(e)}")

if __name__ == "__main__":
    input_path = input("请输入包含Batch IDs的JSONL文件路径: ")
    if os.path.exists(input_path):
        process_batches(input_path)
        print(f"处理完成，结果已保存到 batch_status_output.jsonl")
    else:
        print("错误：输入文件不存在")