#!/usr/bin/env python3
"""Written for https://paraloncloud.com/resources/semantic-search-open-embeddings-api
Needs: pip install numpy scikit-learn requests; PRL_KEY=prlc_... in the environment;
support_tickets_990.jsonl and queries_300.json from the same /files/guides/ folder.

Two production recipes on top of the raw embedding search, measured on the
same 300 multilingual queries: (1) hybrid retrieval, reciprocal-rank fusion
of the embedding ranking and the TF-IDF ranking; (2) an LLM reranker, the 27B
chat model ordering the embedding top-10 for the query. Also the hybrid on the
leave-one-out task."""
import json, os, sys, time, requests, numpy as np, threading, queue
from collections import defaultdict
from sklearn.feature_extraction.text import TfidfVectorizer
EMB="https://paraloncloud.com/v1/embeddings"; CHAT="https://paraloncloud.com/v1/chat/completions"; KEY=os.environ["PRL_KEY"]
rows=[json.loads(l) for l in open("support_tickets_990.jsonl")]; docs=[r["customer_message"] for r in rows]; intents=[r["intent"] for r in rows]
qs=json.load(open("queries.json")); qtexts=[q["query"] for q in qs]; qint=[q["intent"] for q in qs]
def embed(texts):
    out=[]
    for i in range(0,len(texts),64):
        r=requests.post(EMB,headers={"Authorization":f"Bearer {KEY}"},json={"model":"multilingual-e5-small","input":texts[i:i+64]},timeout=120); r.raise_for_status()
        out.extend(e["embedding"] for e in sorted(r.json()["data"],key=lambda e:e["index"]))
    m=np.array(out,dtype=np.float32); return m/np.linalg.norm(m,axis=1,keepdims=True)
P=embed(["passage: "+d for d in docs]); QE=embed(["query: "+t for t in qtexts]); S=QE@P.T
tf=TfidfVectorizer(sublinear_tf=True,ngram_range=(1,2)).fit(docs); T=tf.transform(docs); TQ=(tf.transform(qtexts)@T.T).toarray()
def rrf(*rankings,k=60):
    sc=defaultdict(float)
    for r in rankings:
        for pos,j in enumerate(r): sc[j]+=1/(k+pos+1)
    return [j for j,_ in sorted(sc.items(),key=lambda kv:-kv[1])]
def p15(tops, i): return (intents[tops[0]]==qint[i], sum(intents[j]==qint[i] for j in tops[:5])/5)
res={}
# hybrid on multilingual queries
by=defaultdict(lambda:[0,0.0,0])
for i in range(len(qs)):
    fused=rrf(list(np.argsort(-S[i])[:50]), list(np.argsort(-TQ[i])[:50])); a,b=p15(fused,i); l=qs[i]["lang"]; by[l][0]+=a; by[l][1]+=b; by[l][2]+=1
res["B_hybrid_by_lang"]={l:{"p1":round(100*v[0]/v[2],1),"p5":round(100*v[1]/v[2],1)} for l,v in by.items()}
# hybrid on leave-one-out
Q=embed(["query: "+d for d in docs]); SA=Q@P.T; TA=(T@T.T).toarray(); h1=h5=0
for i in range(len(docs)):
    sa=SA[i].copy(); sa[i]=-1; ta=TA[i].copy(); ta[i]=-1
    fused=rrf(list(np.argsort(-sa)[:50]), list(np.argsort(-ta)[:50])); h1+=intents[fused[0]]==intents[i]; h5+=sum(intents[j]==intents[i] for j in fused[:5])/5
res["A_hybrid"]={"p1":round(100*h1/len(docs),1),"p5":round(100*h5/len(docs),1)}
# LLM rerank of embedding top-10, multilingual queries
schema={"type":"object","properties":{"ranked":{"type":"array","minItems":3,"maxItems":3,"items":{"type":"integer","minimum":1,"maximum":10}}},"required":["ranked"]}
lock=threading.Lock(); stats={"in":0,"out":0,"err":0,"lat":[]}; picks={}
def one(i):
    cand=list(np.argsort(-S[i])[:10])
    listing="\n".join(f"[{n+1}] {docs[j][:300]}" for n,j in enumerate(cand))
    body={"model":"qwen3.8-27b","messages":[{"role":"system","content":"You rank customer-support tickets by how well they match a search query, which may be in another language than the tickets. Return the numbers of the 3 best matches, best first."},
          {"role":"user","content":f"Query: {qtexts[i]}\n\nTickets:\n{listing}"}],
          "response_format":{"type":"json_schema","json_schema":{"name":"rank","schema":schema}},"chat_template_kwargs":{"enable_thinking":False},"temperature":0,"max_tokens":40}
    t=time.time()
    try:
        r=requests.post(CHAT,headers={"Authorization":f"Bearer {KEY}"},json=body,timeout=120); r.raise_for_status(); d=r.json()
        ranked=json.loads(d["choices"][0]["message"]["content"])["ranked"]
        with lock: stats["in"]+=d["usage"]["prompt_tokens"]; stats["out"]+=d["usage"]["completion_tokens"]; stats["lat"].append(time.time()-t); picks[i]=[cand[n-1] for n in ranked]+[j for j in cand if j not in [cand[n-1] for n in ranked]]
    except Exception as e:
        with lock: stats["err"]+=1
        print("ERR",e,file=sys.stderr)
q=queue.Queue(); [q.put(i) for i in range(len(qs))]
def worker():
    while True:
        try: i=q.get_nowait()
        except queue.Empty: return
        one(i); time.sleep(0.2)
th=[threading.Thread(target=worker) for _ in range(12)]; [t.start() for t in th]; [t.join() for t in th]
by=defaultdict(lambda:[0,0.0,0]); byraw=defaultdict(lambda:[0,0])
for i in picks:
    a,b=p15(picks[i],i); l=qs[i]["lang"]; by[l][0]+=a; by[l][1]+=b; by[l][2]+=1
    byraw[l][0]+= intents[int(np.argmax(S[i]))]==qint[i]; byraw[l][1]+=1
res["B_rerank_by_lang"]={l:{"n":v[2],"p1":round(100*v[0]/v[2],1),"p5":round(100*v[1]/v[2],1),"raw_p1":round(100*byraw[l][0]/byraw[l][1],1)} for l,v in by.items()}
lat=sorted(stats["lat"]); res["rerank_cost"]={"queries":len(picks),"errors":stats["err"],"tokens_in_per_query":round(stats["in"]/max(1,len(picks))),"tokens_out_per_query":round(stats["out"]/max(1,len(picks)),1),
    "usd_per_1k_queries":round((stats["in"]*0.12+stats["out"]*1.70)/1e6/max(1,len(picks))*1000,3),"latency_p50_s":round(lat[len(lat)//2],2) if lat else None}
print(json.dumps(res,indent=1)); json.dump(res,open("results_hybrid_rerank.json","w"),indent=1)