File size: 3,663 Bytes
aaf1c39 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """Shared helpers: project paths + OpenRouter chat calls."""
import json, os, re, time, urllib.request
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data"
VECTORS = ROOT / "vectors"
RESULTS = ROOT / "results"
# Target model for the recipe. Switch models with CAPVEC_MODEL (name under models/ or a path),
# so the same pipeline runs on soyuz-4B and Qwen3.5-9B without code edits.
_model = os.environ.get("CAPVEC_MODEL", "soyuz_merged")
MODEL_PATH = _model if os.path.isabs(_model) else str(ROOT / "models" / _model)
# Per-model artifact namespacing (vectors/data of different models must not collide).
MODEL_TAG = os.environ.get("CAPVEC_MODEL_TAG", Path(MODEL_PATH).name)
OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
def openrouter_key() -> str:
# project-local override first (funded key), then the shared env file
for p in (ROOT / ".env.local", Path("/home/alexw/.tbench.env")):
if p.exists():
m = re.search(r"OPENROUTER_API_KEY=(\S+)", p.read_text())
if m:
return m.group(1)
raise RuntimeError("no OPENROUTER_API_KEY found")
def oai_chat(messages, model, base_url=OPENROUTER_URL, api_key=None,
temperature=0.1, max_tokens=4096, retries=3, timeout=120):
"""Call any OpenAI-compatible /chat/completions endpoint (OpenRouter or local vLLM)."""
body = json.dumps({
"model": model,
"messages": messages,
"temperature": temperature,
"max_tokens": max_tokens,
}).encode()
headers = {"Content-Type": "application/json", "User-Agent": "curl/8.5.0"}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
req = urllib.request.Request(base_url, data=body, headers=headers)
last_err = None
for attempt in range(retries):
try:
with urllib.request.urlopen(req, timeout=timeout) as r:
d = json.load(r)
content = d["choices"][0]["message"]["content"]
if content is None:
raise RuntimeError(f"null content, finish_reason={d['choices'][0].get('finish_reason')}")
return content
except Exception as e:
last_err = e
time.sleep(2 * (attempt + 1))
raise last_err
# Backwards-compat wrapper (OpenRouter).
def openrouter_chat(messages, model, temperature=0.1, max_tokens=4096, retries=3, timeout=120):
return oai_chat(messages, model, base_url=OPENROUTER_URL, api_key=openrouter_key(),
temperature=temperature, max_tokens=max_tokens, retries=retries, timeout=timeout)
# Local judge/generator endpoint (Qwen3-8B on vLLM). Append /no_think to disable Qwen3 thinking.
LOCAL_JUDGE_URL = "http://localhost:30008/v1/chat/completions"
LOCAL_JUDGE_MODEL = "qwen3-8b"
def local_chat(messages, temperature=0.1, max_tokens=4096, no_think=True, retries=3, timeout=180):
msgs = [dict(m) for m in messages]
if no_think and msgs:
msgs[-1]["content"] = msgs[-1]["content"] + " /no_think"
return oai_chat(msgs, LOCAL_JUDGE_MODEL, base_url=LOCAL_JUDGE_URL, api_key="dummy",
temperature=temperature, max_tokens=max_tokens, retries=retries, timeout=timeout)
def extract_json(text: str):
"""Pull the first JSON object or array out of a model response (handles ```json fences)."""
fence = re.search(r"```(?:json)?\s*(.*?)```", text, flags=re.S)
if fence:
text = fence.group(1)
start = min([i for i in (text.find("{"), text.find("[")) if i >= 0], default=-1)
if start < 0:
raise ValueError(f"no JSON found in: {text[:200]}")
return json.loads(text[start:])
|