93 lines
4.0 KiB
Python
93 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""Astraea model benchmark — feed identical retrieved facts to each candidate model
|
|
and compare grounded, practical legal answers. Picks the best RAG answer lane.
|
|
"""
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
|
|
sys.path.insert(0, "/opt/astraea")
|
|
import rag # reuse the retrieval pipeline (embed + cosine top-k)
|
|
|
|
OLLAMA = "http://10.30.20.29:11434"
|
|
OUT = "/opt/astraea/benchmark-results.json"
|
|
|
|
MODELS = [
|
|
"granite4.2:latest",
|
|
"qwen3.8fast:latest",
|
|
"ornith-1.5:9b",
|
|
"gemma4:26b-a4b-it-q4_K_M",
|
|
"nemotron-3.5-lightning:30b-a3b-q4_K_M",
|
|
]
|
|
|
|
# (agent_id, question) — each has a real answer in the current index.
|
|
QUESTIONS = [
|
|
("property", "My spouse bought a house before we married, but we paid the mortgage from our joint account for 15 years. When we divorce, is the house community or separate property?"),
|
|
("maintenance", "We're divorcing after a 22-year marriage. I stayed home to raise our kids and have no income. Am I entitled to spousal maintenance, and what factors does the court weigh?"),
|
|
("divorce", "What factors does a Washington court consider when deciding how to divide our property and debts, and can one spouse get more than half?"),
|
|
("divorce", "I just moved to Washington three weeks ago and want to file for divorce from my spouse who lives in Oregon. Can I file here, and is there a waiting period?"),
|
|
]
|
|
|
|
SYSTEM = (
|
|
"You are a licensed Washington State family-law attorney giving a real client direct, "
|
|
"practical legal advice. Base every legal claim on the reference documents. Cite the RCW "
|
|
"section AND any controlling case law by name. Be specific and actionable: state the rule, "
|
|
"apply it to the client's facts, and tell them what to do next. Do not hedge or refuse to "
|
|
"answer. If the exact scenario is not in the documents, give the closest governing principle "
|
|
"and note it needs attorney review."
|
|
)
|
|
|
|
|
|
def chat(model, system, user):
|
|
payload = {
|
|
"model": model,
|
|
"messages": [{"role": "system", "content": system},
|
|
{"role": "user", "content": user}],
|
|
"stream": False,
|
|
"think": False,
|
|
"options": {"temperature": 0.2, "num_predict": 1400, "num_ctx": 16384},
|
|
}
|
|
req = urllib.request.Request(f"{OLLAMA}/api/chat", data=json.dumps(payload).encode(),
|
|
headers={"Content-Type": "application/json"})
|
|
op = urllib.request.build_opener(urllib.request.ProxyHandler({}))
|
|
with op.open(req, timeout=300) as r:
|
|
d = json.loads(r.read().decode("utf-8"))
|
|
m = d.get("message", {})
|
|
return (m.get("content") or m.get("thinking") or "").strip()
|
|
|
|
|
|
def main():
|
|
if not rag._INDEX:
|
|
print("loading index from cache...", flush=True)
|
|
rag._INDEX.extend(rag.build_index())
|
|
print(f"index: {len(rag._INDEX)} chunks", flush=True)
|
|
results = []
|
|
for agent_id, q in QUESTIONS:
|
|
top = rag.retrieve(q, agent_id, top_k=5)
|
|
ctx = "\n\n".join(
|
|
f"[{i+1}] ({c['source']} — {c['title']}):\n{c['text']}"
|
|
for i, (s, c) in enumerate(top)
|
|
)
|
|
print(f"\n=== Q: {q[:70]}... ({len(top)} chunks) ===", flush=True)
|
|
for model in MODELS:
|
|
t0 = time.time()
|
|
try:
|
|
ans = chat(model, SYSTEM, f"REFERENCE DOCUMENTS:\n\n{ctx}\n\nQUESTION: {q}")
|
|
dt = round(time.time() - t0, 1)
|
|
print(f" {model}: {len(ans)} chars in {dt}s", flush=True)
|
|
results.append({"question": q, "agent": agent_id, "model": model,
|
|
"answer": ans, "latency_s": dt,
|
|
"sources": [c["source"] for s, c in top]})
|
|
except Exception as e:
|
|
print(f" {model}: ERROR {e}", flush=True)
|
|
results.append({"question": q, "agent": agent_id, "model": model, "error": str(e)})
|
|
json.dump(results, open(OUT, "w"), indent=2)
|
|
print(f"\n=== done -> {OUT} ===")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|