| """Run the agentic controller over a query file -> run.tsv (final ranked docs), for nDCG scoring. |
| Two phases: (1) batch-plan all queries in parallel (DeepSeek, cached per qid so ablation runs reuse |
| the SAME plans), (2) assemble each query (parallel fetch + rerank + rank) sequentially. |
| |
| Env: THEMIS_DATA, THEMIS_STATUTE, THEMIS_QFILE, OUT, PLANCACHE, THEMIS_ENABLED=all|vector,authority,... |
| """ |
| import os, sys, json, time |
| import concurrent.futures as cf |
| import requests |
| sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "scripts")) |
| from tools import Corpus |
| import agent as A |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| def _load_env(p): |
| for l in open(p): |
| l = l.strip() |
| if l and not l.startswith("#") and "=" in l: |
| k, v = l.split("=", 1); os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'")) |
| _load_env(os.path.join(HERE, "..", "scripts", ".env")) |
| HDR = {"Authorization": f"Bearer {os.environ['DEEPSEEK_API_KEY']}", "Content-Type": "application/json"} |
|
|
| def llm_fn(msgs): |
| for _ in range(3): |
| try: |
| r = requests.post("https://api.deepseek.com/chat/completions", headers=HDR, timeout=60, |
| json={"model": "deepseek-chat", "temperature": 0, "max_tokens": 400, "messages": msgs}) |
| if r.status_code == 200: return r.json()["choices"][0]["message"]["content"] |
| except Exception: time.sleep(2) |
| return "{}" |
|
|
| QFILE = os.environ.get("THEMIS_QFILE", "authority_queries.tsv") |
| OUT = os.environ.get("OUT", "agent_run.tsv") |
| PLANCACHE = os.environ.get("PLANCACHE", "plan_" + os.path.basename(QFILE).split(".")[0] + ".json") |
| ENABLED = A.ALL_TOOLS if os.environ.get("THEMIS_ENABLED", "all") == "all" else set(os.environ["THEMIS_ENABLED"].split(",")) |
| ALPHA = float(os.environ.get("ALPHA", "0.3")) |
|
|
| rows = [] |
| for l in open(QFILE, encoding="utf-8"): |
| qid, it, t = l.rstrip("\n").split("\t", 2); rows.append((qid, t)) |
|
|
| |
| plans = json.load(open(PLANCACHE)) if os.path.exists(PLANCACHE) else {} |
| todo = [(qid, t) for qid, t in rows if qid not in plans] |
| if todo: |
| print(f"planning {len(todo)} queries ...", flush=True); t0 = time.time() |
| with cf.ThreadPoolExecutor(max_workers=24) as ex: |
| futs = {ex.submit(A.plan, t, llm_fn): qid for qid, t in todo} |
| done = 0 |
| for f in cf.as_completed(futs): |
| plans[futs[f]] = f.result(); done += 1 |
| if done % 50 == 0: json.dump(plans, open(PLANCACHE, "w")); print(f" {done}/{len(todo)} {time.time()-t0:.0f}s", flush=True) |
| json.dump(plans, open(PLANCACHE, "w")) |
| print(f"plans done {time.time()-t0:.0f}s", flush=True) |
|
|
| |
| C = Corpus(os.environ.get("THEMIS_DATA", "."), os.environ.get("THEMIS_STATUTE", "."), device=os.environ.get("THEMIS_DEVICE", "cpu")) |
| print(f"assembling {len(rows)} queries (enabled={sorted(ENABLED)}) ...", flush=True) |
| t0 = time.time() |
| with open(OUT, "w", encoding="utf-8") as f: |
| for i, (qid, t) in enumerate(rows): |
| ranked, info = A.assemble(C, t, plans[qid], enabled=ENABLED, alpha=ALPHA) |
| for rank, d in enumerate(ranked, 1): f.write(f"{qid}\t{rank}\t{d}\n") |
| if (i + 1) % 30 == 0: print(f" {i+1}/{len(rows)} {time.time()-t0:.0f}s", flush=True) |
| print(f"DONE {len(rows)} -> {OUT} in {time.time()-t0:.0f}s", flush=True) |
|
|
