#!/usr/bin/env python3
"""Written for https://paraloncloud.com/resources/batch-api-open-models
Needs: pip install openai; PRL_KEY=prlc_... in the environment; support_tickets_990.jsonl from
https://paraloncloud.com/files/guides/llm-as-a-judge-open-model-api/support_tickets_990.jsonl next to it.

Batch API on open models, measured with the official OpenAI Python SDK.

1. 300 support tickets classified by intent (chat completions, JSON schema)
2. 100 tickets embedded (embeddings)
3. 20 tickets answered through the Responses API
Each: files.create -> batches.create -> poll -> files.content, timed.
While the chat batch runs, live requests are timed against a baseline taken
just before, to see whether the batch slows real-time traffic.
"""
import json, os, sys, time, statistics, threading
from openai import OpenAI

BASE = "https://paraloncloud.com/v1"; KEY = os.environ["PRL_KEY"]
client = OpenAI(base_url=BASE, api_key=KEY)
CHAT, EMBED = "qwen3.8-27b", "multilingual-e5-small"
rows = [json.loads(l) for l in open(os.environ.get("DATA", "support_tickets_990.jsonl"))]
INTENTS = sorted({r["intent"] for r in rows})
SCHEMA = {"type": "object", "properties": {"intent": {"type": "string", "enum": INTENTS}}, "required": ["intent"]}
res = {}

def write_jsonl(path, lines):
    with open(path, "w") as f:
        for l in lines:
            f.write(json.dumps(l) + "\n")

def run_batch(name, endpoint, lines):
    path = f"{name}.jsonl"; write_jsonl(path, lines)
    t0 = time.time()
    f = client.files.create(file=open(path, "rb"), purpose="batch")
    b = client.batches.create(input_file_id=f.id, endpoint=endpoint, completion_window="24h", metadata={"bench": name})
    t_created = time.time()
    progress = []
    while b.status not in ("completed", "failed", "expired", "cancelled"):
        time.sleep(2)
        b = client.batches.retrieve(b.id)
        progress.append((round(time.time() - t_created, 1), b.request_counts.completed))
    t_done = time.time()
    out = client.files.content(b.output_file_id).text if b.output_file_id else ""
    err = client.files.content(b.error_file_id).text if b.error_file_id else ""
    results = {}
    for line in out.splitlines():
        o = json.loads(line); results[o["custom_id"]] = o
    info = {"batch_id": b.id, "status": b.status, "lines": len(lines), "completed": b.request_counts.completed,
            "failed": b.request_counts.failed, "error_lines": len(err.splitlines()), "upload_and_create_s": round(t_created - t0, 2),
            "run_s": round(t_done - t_created, 1), "per_min": round(len(lines) / max(1e-9, (t_done - t_created)) * 60, 1)}
    print(name, json.dumps(info), file=sys.stderr)
    return info, results, progress

def live_latency(n, tag):
    lat = []
    for i in range(n):
        t = time.time()
        client.chat.completions.create(model=CHAT, messages=[{"role": "user", "content": "Name the capital of France in one word."}],
                                       max_tokens=10, temperature=0, extra_body={"chat_template_kwargs": {"enable_thinking": False}})
        lat.append(time.time() - t)
        time.sleep(1)
    return {"n": n, "p50_s": round(statistics.median(lat), 3), "max_s": round(max(lat), 3)}

# ---- baseline live latency
res["live_before"] = live_latency(15, "before")

# ---- 1. chat classification, 300 lines, with live latency sampled during the run
chat_rows = rows[:300]
chat_lines = [{"custom_id": f"t{i}", "method": "POST", "url": "/v1/chat/completions", "body": {
    "model": CHAT, "temperature": 0, "max_tokens": 30, "chat_template_kwargs": {"enable_thinking": False},
    "response_format": {"type": "json_schema", "json_schema": {"name": "intent", "schema": SCHEMA}},
    "messages": [{"role": "system", "content": "Classify the customer-support message into one intent."},
                 {"role": "user", "content": r["customer_message"]}]}} for i, r in enumerate(chat_rows)]
during = {}
def sample_during():
    time.sleep(8)
    during["v"] = live_latency(15, "during")
th = threading.Thread(target=sample_during); th.start()
info, results, progress = run_batch("chat300", "/v1/chat/completions", chat_lines)
th.join()
correct = 0
for i, r in enumerate(chat_rows):
    o = results.get(f"t{i}")
    try:
        got = json.loads(o["response"]["body"]["choices"][0]["message"]["content"])["intent"]
        correct += got == r["intent"]
    except Exception:
        pass
info["accuracy_pct"] = round(100 * correct / len(chat_rows), 1)
res["chat300"] = info; res["chat300_progress"] = progress[::5]; res["live_during"] = during.get("v")

# ---- 2. embeddings, 100 lines
emb_lines = [{"custom_id": f"e{i}", "method": "POST", "url": "/v1/embeddings",
              "body": {"model": EMBED, "input": "passage: " + r["customer_message"]}} for i, r in enumerate(rows[300:400])]
info, results, _ = run_batch("embed100", "/v1/embeddings", emb_lines)
dims = set()
for o in results.values():
    dims.add(len(o["response"]["body"]["data"][0]["embedding"]))
info["dims"] = sorted(dims); res["embed100"] = info

# ---- 3. responses, 20 lines
resp_lines = [{"custom_id": f"r{i}", "method": "POST", "url": "/v1/responses", "body": {
    "model": CHAT, "reasoning": {"effort": "low"}, "max_output_tokens": 120,
    "instructions": "You are a support agent. Draft a two-sentence reply.", "input": r["customer_message"]}} for i, r in enumerate(rows[400:420])]
info, results, _ = run_batch("responses20", "/v1/responses", resp_lines)
ok = sum(1 for o in results.values() if o["response"]["body"].get("status") in ("completed", "incomplete") and o["response"]["body"].get("output"))
info["responses_with_output"] = ok; res["responses20"] = info

json.dump(res, open("results.json", "w"), indent=1)
print(json.dumps(res, indent=1))