scripts/common.py

"""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:])