#!/usr/bin/env python3
"""RAG answers checked sentence by sentence with /v1/classify, scored on
planted hallucinations so recall and false-alarm rate are exact.
  PRL_KEY=prlc_... DOCS=./docs python3 rag_check.py
Steps: chunk the docs -> embed (multilingual-e5-small) -> generate questions
from random chunks (qwen3.8-27b) -> retrieve top-3 -> answer (qwen3.8-27b)
-> plant one wrong sentence in half the answers -> classify every sentence
against the retrieved passages (openjev-4b) -> report. Optionally judge the
same sentences with the 27B (JUDGE=1) for comparison.
Written for https://paraloncloud.com/resources/rag-hallucination-check-classify-api
"""
import json, os, re, glob, random, time, sys, math, threading, queue, requests
BASE="https://paraloncloud.com/v1"; KEY=os.environ["PRL_KEY"]; H={"Authorization":"Bearer "+KEY}
DOCS=os.environ.get("DOCS","./docs"); N_Q=int(os.environ.get("N_Q","50")); SEED=int(os.environ.get("SEED","7")); JUDGE=os.environ.get("JUDGE","1")=="1"
OUT=os.environ.get("OUT","rag_check_results.json"); rnd=random.Random(SEED)
sess=requests.Session()
def post(path,body,timeout=180):
    for a in range(4):
        try:
            r=sess.post(BASE+path,headers=H,json=body,timeout=timeout)
            if r.status_code==200: return r.json()
            print("HTTP",r.status_code,r.text[:120],file=sys.stderr)
        except Exception as e: print("EXC",type(e).__name__,file=sys.stderr)
        time.sleep(3)
    raise SystemExit("gave up on "+path)
# 1. chunks
chunks=[]
for f in sorted(glob.glob(f"{DOCS}/**/*.md",recursive=True)):
    t=open(f).read(); t=re.sub(r"^---.*?---\s*","",t,flags=re.S)
    parts=re.split(r"\n(?=#{1,3} )",t)
    for p in parts:
        p=re.sub(r"```.*?```","",p,flags=re.S); p=re.sub(r"\n{2,}","\n",p).strip()
        words=p.split()
        if 60<=len(words)<=450: chunks.append({"src":os.path.relpath(f,DOCS),"text":" ".join(words)})
print(f"{len(chunks)} chunks from {DOCS}",file=sys.stderr)
# 2. embed
def embed(texts,prefix):
    vecs=[]
    for i in range(0,len(texts),64):
        d=post("/embeddings",{"model":"multilingual-e5-small","input":[prefix+t for t in texts[i:i+64]]})
        vecs+= [e["embedding"] for e in sorted(d["data"],key=lambda e:e["index"])]
    return vecs
cv=embed([c["text"] for c in chunks],"passage: ")
# 3. questions from random chunks
per=max(1,-(-N_Q//len(chunks))); pick=rnd.sample(range(len(chunks)),min(N_Q,len(chunks))); questions=[]
for ci in pick:
    d=post("/chat/completions",{"model":"qwen3.8-27b","temperature":0,"max_tokens":40*per+20,"chat_template_kwargs":{"enable_thinking":False},
        "messages":[{"role":"system","content":f"Write {per} different short question(s) a user of this product would ask that the passage answers, one per line, no numbering. Output only the questions."},{"role":"user","content":chunks[ci]["text"]}]})
    for line in d["choices"][0]["message"]["content"].strip().split("\n")[:per]:
        line=line.strip().lstrip("-*0123456789. ").strip()
        if line: questions.append({"q":line,"src_chunk":ci})
questions=questions[:N_Q]
qv=embed([q["q"] for q in questions],"query: ")
def cos(a,b): return sum(x*y for x,y in zip(a,b))/(math.sqrt(sum(x*x for x in a))*math.sqrt(sum(y*y for y in b)))
# 4. retrieve + answer
hit1=0; answers=[]
for q,v in zip(questions,qv):
    top=sorted(range(len(chunks)),key=lambda i:-cos(v,cv[i]))[:3]; hit1+=(top[0]==q["src_chunk"])
    ctx="\n\n".join(f"[{k+1}] {chunks[i]['text']}" for k,i in enumerate(top))
    d=post("/chat/completions",{"model":"qwen3.8-27b","temperature":0,"max_tokens":220,"chat_template_kwargs":{"enable_thinking":False},
        "messages":[{"role":"system","content":"Answer the question in 3 to 5 plain sentences using only the passages. No bullet points, no citations, no headings."},{"role":"user","content":f"Passages:\n{ctx}\n\nQuestion: {q['q']}"}]})
    answers.append({**q,"top":top,"answer":d["choices"][0]["message"]["content"].strip()})
print(f"retrieval top-1 hit: {hit1}/{len(questions)}",file=sys.stderr)
# 5. split + plant
def sentences(t): return [s.strip() for s in re.split(r"(?<=[.!?])\s+(?=[A-Z0-9])",re.sub(r"\s+"," ",t)) if len(s.strip())>20]
CLAIMS=["This feature is included free of charge on every enterprise plan.","Refunds are processed automatically within 24 hours.","The limit is lifted for accounts older than one year.",
 "Support is available by phone around the clock.","A dedicated GPU is reserved for each account by default.","This setting cannot be changed once the account is created."]
def plant(sent):
    m=re.search(r"(?<![A-Za-z\d#.-])\d+(?:[.,]\d+)?(?![A-Za-z\d])",sent)  # a number, not part of WSL2 / #48291 / v0.29
    if m and float(m.group(0).replace(",",""))!=0 and rnd.random()<0.6:   # 0 x 10 is still 0: not an error
        num=m.group(0).replace(",",""); new=str(int(float(num)*10)) if float(num)==int(float(num)) else f"{float(num)*10:g}"
        return sent[:m.start()]+new+sent[m.end():],"number"
    if re.search(r"\b(is|are|can|will|does|has)\b",sent) and " not " not in sent and "n't" not in sent and rnd.random()<0.5:
        return re.sub(r"\b(is|are|can|will|does|has)\b",lambda x:x.group(0)+" not",sent,count=1),"negation"
    return sent,"replace"
rows=[]
for ai,a in enumerate(answers):
    ss=sentences(a["answer"]); 
    if len(ss)<2: continue
    planted=None
    if ai%2==1:
        j=rnd.randrange(len(ss)); new,kind=plant(ss[j])
        if kind=="replace": new=rnd.choice(CLAIMS)
        ss[j]=new; planted=(j,kind)
    for j,s in enumerate(ss): rows.append({"ai":ai,"j":j,"sent":s,"planted":bool(planted and planted[0]==j),"kind":planted[1] if planted and planted[0]==j else None,"top":a["top"]})
print(f"{len(rows)} sentences, {sum(r['planted'] for r in rows)} planted",file=sys.stderr)
# 6. classify: each sentence vs each of its 3 passages; support = max entailment
pairs=[(ri,ci) for ri,r in enumerate(rows) for ci in r["top"]]; ent={}; t0=time.time(); tok=0
for b in range(0,len(pairs),128):
    chunk=pairs[b:b+128]
    d=post("/classify",{"model":"openjev-4b","pairs":[{"premise":chunks[ci]["text"],"hypothesis":rows[ri]["sent"]} for ri,ci in chunk]})
    tok+=d["usage"]["prompt_tokens"]
    for (ri,ci),it in zip(chunk,d["data"]): ent[(ri,ci)]=it["probs"]
cls_wall=time.time()-t0
for ri,r in enumerate(rows):
    bc=max(r["top"],key=lambda ci:ent[(ri,ci)][1]); p=ent[(ri,bc)]; r["support"]=p[1]
    r["label"]=("contradiction","entailment","neutral")[max(range(3),key=lambda j:p[j])]  # verdict of the passage that supports it most
def metrics(flag):
    tp=sum(1 for r in rows if r["planted"] and flag(r)); fn=sum(1 for r in rows if r["planted"] and not flag(r))
    fp=sum(1 for r in rows if not r["planted"] and flag(r)); tn=sum(1 for r in rows if not r["planted"] and not flag(r))
    return {"planted_caught_pct":round(100*tp/max(1,tp+fn),1),"clean_flagged_pct":round(100*fp/max(1,fp+tn),1),"tp":tp,"fn":fn,"fp":fp,"tn":tn}
res={"chunks":len(chunks),"questions":len(questions),"retrieval_top1_hit":hit1,"sentences":len(rows),"planted":sum(r["planted"] for r in rows),
 "classify":{"pairs":len(pairs),"tokens":tok,"cost_usd":round(tok*0.05/1e6,4),"wall_s":round(cls_wall,1),
   "by_threshold":{str(t):metrics(lambda r,t=t:r["support"]<t) for t in (0.3,0.5,0.7)},
   "by_label":metrics(lambda r:r["label"]!="entailment"),
   "by_kind":{k:round(100*sum(1 for r in rows if r["kind"]==k and r["support"]<0.5)/max(1,sum(1 for r in rows if r["kind"]==k)),1) for k in ("number","negation","replace")}}}
# 7. optional: 27B judge per sentence
if JUDGE:
    lock=threading.Lock(); jud={}; jt=[0,0,0]; t0=time.time(); qq=queue.Queue(); [qq.put(ri) for ri in range(len(rows))]
    def w():
        while True:
            try: ri=qq.get_nowait()
            except queue.Empty: return
            r=rows[ri]; ctx="\n\n".join(chunks[ci]["text"] for ci in r["top"])
            try:
                d=post("/chat/completions",{"model":"qwen3.8-27b","temperature":0,"max_tokens":8,"chat_template_kwargs":{"enable_thinking":False},
                    "messages":[{"role":"system","content":"You check a sentence against source passages. Answer exactly one word: supported if the passages state or clearly imply the sentence, unsupported otherwise."},{"role":"user","content":f"Passages:\n{ctx}\n\nSentence: {r['sent']}"}]})
                a=d["choices"][0]["message"]["content"].strip().lower()
                with lock: jud[ri]=("unsupported" in a); jt[0]+=d["usage"]["prompt_tokens"]; jt[1]+=d["usage"]["completion_tokens"]
            except SystemExit:
                with lock: jt[2]+=1
    th=[threading.Thread(target=w) for _ in range(8)]; [t.start() for t in th]; [t.join() for t in th]
    res["judge_27b"]={**metrics(lambda r:jud.get(rows.index(r),False)),"tokens_in":jt[0],"tokens_out":jt[1],"failed":jt[2],"cost_usd":round((jt[0]*0.12+jt[1]*1.7)/1e6,4),"wall_s":round(time.time()-t0,1)}
    for ri,r in enumerate(rows): r["judge_flag"]=jud.get(ri)
print(json.dumps(res,indent=1)); json.dump({"summary":res,"rows":rows,"answers":answers,"questions":questions},open(OUT,"w"),indent=1)