#!/usr/bin/env python3
"""LLM-as-a-judge with an open 27B model, scored against known truth.
Label verification: a support ticket plus a proposed intent label; half the
labels are the true one, half are a deliberately wrong one from the same
taxonomy. The judge answers correct / incorrect with a one-line reason.
Recall on wrong labels and acceptance of right ones are exact because the
errors were planted.

  PRL_KEY=prlc_... RUBRIC=1 DATA=support_tickets_990.jsonl RPM=20 IN_FLIGHT=2 python3 judge_labels.py

RUBRIC=1 adds a one-line definition per intent and the single-best-fit rule;
RUBRIC=0 gives the judge only the task. Point DATA at any JSONL with
"customer_message" and "intent" fields and edit INTENTS/DEFS for your taxonomy.

Written for https://paraloncloud.com/resources/llm-as-a-judge-open-model-api
"""
import json, os, sys, time, random, threading, queue, requests
from collections import Counter
API="https://paraloncloud.com/v1/chat/completions"; MODEL=os.environ.get("MODEL","qwen3.8-27b"); KEY=os.environ["PRL_KEY"]
DATA=os.environ.get("DATA","support_tickets_990.jsonl"); N=int(os.environ.get("N","990")); IN_FLIGHT=int(os.environ.get("IN_FLIGHT","2")); RPM=int(os.environ.get("RPM","20"))   # free key: 20 rpm
RUBRIC=os.environ.get("RUBRIC","1")=="1"   # with vs without a written rubric
PRICE_IN,PRICE_OUT=0.12,1.70   # USD per 1M tokens, qwen3.8-27b list price
INTENTS=["order_status","cancel_order","return_request","refund_status","damaged_item","wrong_item","payment_failed","discount_code","product_question","shipping_address_change","account_login","complaint"]
DEFS={"order_status":"asks where an order is or when it arrives","cancel_order":"wants to cancel an order not yet delivered","return_request":"wants to send an item back",
 "refund_status":"asks about money owed back for a return or cancellation already agreed","damaged_item":"item arrived broken or defective","wrong_item":"received a different item than ordered",
 "payment_failed":"card or payment did not go through","discount_code":"a promo code does not work or how to use one","product_question":"asks about a product before or after buying, not a problem with it",
 "shipping_address_change":"wants to change where an order ships","account_login":"cannot sign in, password, account access","complaint":"expresses dissatisfaction with service or the company as a whole"}
rows=[json.loads(l) for l in open(DATA)][:N]; rnd=random.Random(21)
cases=[]
for i,r in enumerate(rows):
    if i%2==0: cases.append({"msg":r["customer_message"],"label":r["intent"],"truth":True})
    else:
        wrong=rnd.choice([x for x in INTENTS if x!=r["intent"]]); cases.append({"msg":r["customer_message"],"label":wrong,"truth":False,"real":r["intent"]})
SYSTEM=("You are a strict quality reviewer for a customer-support dataset. Decide whether the proposed intent label is the correct one for the message. "
        +("Definitions: "+"; ".join(f"{k}: {v}" for k,v in DEFS.items())+". A label is correct only if it is the single best fit; if another label fits clearly better, it is incorrect. " if RUBRIC else "")
        +"Answer via the schema.")
SCHEMA={"type":"object","properties":{"verdict":{"type":"string","enum":["correct","incorrect"]},"reason":{"type":"string","maxLength":160}},"required":["verdict","reason"]}
lock=threading.Lock(); stats={"req":0,"err":0,"in":0,"out":0,"lat":[]}; out=[]
def one(c):
    user=f"Message: {c['msg']}\nProposed intent: {c['label']}"
    body={"model":MODEL,"messages":[{"role":"system","content":SYSTEM},{"role":"user","content":user}],
          "response_format":{"type":"json_schema","json_schema":{"name":"verdict","schema":SCHEMA}},"chat_template_kwargs":{"enable_thinking":False},"temperature":0,"max_tokens":120}
    t0=time.time()
    try: r=requests.post(API,headers={"Authorization":f"Bearer {KEY}"},json=body,timeout=180)
    except Exception as e:
        with lock: stats["req"]+=1; stats["err"]+=1
        return
    lat=time.time()-t0
    with lock: stats["req"]+=1; stats["lat"].append(lat)
    if r.status_code!=200:
        with lock: stats["err"]+=1
        print("HTTP",r.status_code,r.text[:120],file=sys.stderr); return
    d=r.json(); u=d.get("usage",{})
    try: v=json.loads(d["choices"][0]["message"]["content"])
    except Exception as e:
        with lock: stats["err"]+=1
        return
    with lock:
        stats["in"]+=u.get("prompt_tokens",0); stats["out"]+=u.get("completion_tokens",0)
        out.append({**c,"verdict":v["verdict"]=="correct","reason":v.get("reason","")})
q=queue.Queue(); [q.put(c) for c in cases]; start=time.time()
def worker():
    while True:
        try: c=q.get_nowait()
        except queue.Empty: return
        with lock: ahead=stats["req"]+IN_FLIGHT-((time.time()-start)/60*RPM+1)
        if ahead>0: time.sleep(min(5,ahead*60/RPM))
        one(c)
th=[threading.Thread(target=worker) for _ in range(IN_FLIGHT)]; [t.start() for t in th]; [t.join() for t in th]; wall=time.time()-start
tp=sum(1 for o in out if o["truth"] and o["verdict"]); tn=sum(1 for o in out if not o["truth"] and not o["verdict"])
fp=sum(1 for o in out if not o["truth"] and o["verdict"]); fn=sum(1 for o in out if o["truth"] and not o["verdict"])
n=len(out); lat=sorted(stats["lat"]); pct=lambda x: round(lat[min(len(lat)-1,int(len(lat)*x))],2) if lat else None
missed=Counter(f"{o['real']} labelled as {o['label']}" for o in out if not o["truth"] and o["verdict"])
m={"rubric":RUBRIC,"cases":n,"errors":stats["err"],"accuracy_pct":round(100*(tp+tn)/n,2),
   "correct_labels_accepted_pct":round(100*tp/max(1,tp+fn),2),"wrong_labels_caught_pct":round(100*tn/max(1,tn+fp),2),
   "precision_of_reject_pct":round(100*tn/max(1,tn+fn),2),"tp":tp,"tn":tn,"fp":fp,"fn":fn,
   "tokens_in_per_case":round(stats["in"]/n,1),"tokens_out_per_case":round(stats["out"]/n,1),"latency_p50_s":pct(0.5),"latency_p90_s":pct(0.9),
   "wall_s":round(wall,1),"cases_per_hour":round(n/wall*3600),"cost_per_1k_cases_usd":round((stats["in"]*PRICE_IN+stats["out"]*PRICE_OUT)/1e6/n*1000,4),
   "wrong_labels_missed_top":missed.most_common(6),
   "false_rejects_sample":[{"msg":o["msg"][:140],"label":o["label"],"reason":o["reason"]} for o in out if o["truth"] and not o["verdict"]][:8]}
print(json.dumps(m,indent=1)); json.dump({"summary":m,"rows":out},open(f"results_rubric{int(RUBRIC)}.json","w"),indent=1)