使用 Batch API 批量处理任务

import base64
import json
import os
import time
from pathlib import Path

from openai import OpenAI
from openai.types import Batch, FileObject

client = OpenAI(
    api_key=os.environ.get("MOONSHOT_API_KEY"),
    base_url=os.environ.get("MOONSHOT_BASE_URL", "https://api.moonshot.cn/v1"),
)

MODEL = "kimi-k2.6"
PROMPT = "请分类这张图片：风景/人物/美食/建筑/其他"
SYSTEM = "你是一个图片分类助手"


def build_request_base64(custom_id: str, image_path: str) -> dict:
    """方式一：将图片编码为 base64 直接内嵌到 JSONL 中。
    适合小图片，无需额外上传步骤。"""
    with open(image_path, "rb") as f:
        image_data: str = base64.b64encode(f.read()).decode("utf-8")
    return {
        "custom_id": custom_id,
        "method": "POST",
        "url": "/v1/chat/completions",
        "body": {
            "model": MODEL,
            "messages": [
                {"role": "system", "content": SYSTEM},
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_data}"}},
                        {"type": "text", "text": PROMPT},
                    ],
                },
            ],
        },
    }


def build_request_upload(custom_id: str, image_path: str) -> dict:
    """方式二：先上传图片获取 file_id，再通过 ms://<file_id> 引用。
    适合大图片或同一张图片被多个请求复用的场景。"""
    file_object: FileObject = client.files.create(
        file=open(image_path, "rb"),
        purpose="image",
    )
    print(f"图片已上传: {image_path} -> {file_object.id}")
    return {
        "custom_id": custom_id,
        "method": "POST",
        "url": "/v1/chat/completions",
        "body": {
            "model": MODEL,
            "messages": [
                {"role": "system", "content": SYSTEM},
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": f"ms://{file_object.id}"}},
                        {"type": "text", "text": PROMPT},
                    ],
                },
            ],
        },
    }


# ====== 在这里选择构建方式 ======
build_request = build_request_base64  # 或 build_request_upload
# ================================

# 1. 构造输入文件
images: list[str] = ["image1.png", "image2.png", "image3.png"]
requests: list[dict] = [build_request(f"img-{i}", path) for i, path in enumerate(images)]

input_path = Path("image_batch_requests.jsonl")
with input_path.open("w", encoding="utf-8") as f:
    for req in requests:
        f.write(json.dumps(req, ensure_ascii=False) + "\n")

# 2. 上传 JSONL 并创建任务
file_object: FileObject = client.files.create(file=input_path, purpose="batch")
batch: Batch = client.batches.create(
    input_file_id=file_object.id,
    endpoint="/v1/chat/completions",
    completion_window="24h",
)
print(f"任务已创建: {batch.id}")

# 3. 轮询等待完成
while True:
    batch = client.batches.retrieve(batch.id)
    print(f"状态: {batch.status} ({batch.request_counts.completed}/{batch.request_counts.total})")
    if batch.status == "completed":
        break
    elif batch.status in ("failed", "expired", "cancelled"):
        print(f"任务异常终止: {batch.status}")
        exit(1)
    time.sleep(10)

# 4. 处理结果
output = client.files.content(batch.output_file_id)
for line in output.text.strip().split("\n"):
    data: dict = json.loads(line)
    print(f"{data['custom_id']}: {data['response']['body']['choices'][0]['message']['content']}")

import base64
import json
import os
import time
from pathlib import Path

from openai import OpenAI
from openai.types import Batch, FileObject

MODEL = "kimi-k2.6"

client = OpenAI(
    api_key=os.environ.get("MOONSHOT_API_KEY"),
    base_url=os.environ.get("MOONSHOT_BASE_URL", "https://api.moonshot.cn/v1"),
)

PROMPT = "请总结这个视频的主要内容"
SYSTEM = "你是一个视频内容分析助手"


def build_request_base64(custom_id: str, video_path: str) -> dict:
    """方式一：将视频编码为 base64 直接内嵌到 JSONL 中。
    适合小视频，无需额外上传步骤。"""
    with open(video_path, "rb") as f:
        video_data: str = base64.b64encode(f.read()).decode("utf-8")
    return {
        "custom_id": custom_id,
        "method": "POST",
        "url": "/v1/chat/completions",
        "body": {
            "model": MODEL,
            "messages": [
                {"role": "system", "content": SYSTEM},
                {
                    "role": "user",
                    "content": [
                        {"type": "video_url", "video_url": {"url": f"data:video/mp4;base64,{video_data}"}},
                        {"type": "text", "text": PROMPT},
                    ],
                },
            ],
        },
    }


def build_request_upload(custom_id: str, video_path: str) -> dict:
    """方式二：先上传视频获取 file_id，再通过 ms://<file_id> 引用。
    适合大视频或同一个视频被多个请求复用的场景。"""
    file_object: FileObject = client.files.create(
        file=open(video_path, "rb"),
        purpose="video",
    )
    print(f"视频已上传: {video_path} -> {file_object.id}")
    return {
        "custom_id": custom_id,
        "method": "POST",
        "url": "/v1/chat/completions",
        "body": {
            "model": MODEL,
            "messages": [
                {"role": "system", "content": SYSTEM},
                {
                    "role": "user",
                    "content": [
                        {"type": "video_url", "video_url": {"url": f"ms://{file_object.id}"}},
                        {"type": "text", "text": PROMPT},
                    ],
                },
            ],
        },
    }


# ====== 在这里选择构建方式 ======
build_request = build_request_base64  # 或 build_request_upload
# ================================

# 1. 构造输入文件
videos: list[str] = ["video1.mp4", "video2.mp4", "video3.mp4"]
requests: list[dict] = [build_request(f"video-{i}", path) for i, path in enumerate(videos)]

input_path = Path("video_batch_requests.jsonl")
with input_path.open("w", encoding="utf-8") as f:
    for req in requests:
        f.write(json.dumps(req, ensure_ascii=False) + "\n")

# 2. 上传 JSONL 并创建任务
batch_file: FileObject = client.files.create(file=input_path, purpose="batch")
batch: Batch = client.batches.create(
    input_file_id=batch_file.id,
    endpoint="/v1/chat/completions",
    completion_window="24h",
)
print(f"任务已创建: {batch.id}")

# 3. 轮询等待完成
while True:
    batch = client.batches.retrieve(batch.id)
    print(f"状态: {batch.status} ({batch.request_counts.completed}/{batch.request_counts.total})")
    if batch.status == "completed":
        break
    elif batch.status in ("failed", "expired", "cancelled"):
        print(f"任务异常终止: {batch.status}")
        exit(1)
    time.sleep(10)

# 4. 处理结果
output = client.files.content(batch.output_file_id)
for line in output.text.strip().split("\n"):
    data: dict = json.loads(line)
    print(f"{data['custom_id']}: {data['response']['body']['choices'][0]['message']['content']}")