AILang/experiments/2026-05-21-naming-ab/run.py

#!/usr/bin/env python3
"""A/B test: does Qwen3-Coder produce valid Form-A code more reliably
when the spec uses long tag names (apply/lambda/constructor/typecon)
vs. the current short ones (app/lam/ctor/con)?

Setup:
- Two specs, identical except for tag names.
- Same N tasks, fired against each cohort.
- For each call: log raw response, run it through `ail check`
  (renamed outputs go through reverse substitution first), capture
  parser / typecheck / build / run results.

Token budget guard: aborts before exceeding --token-budget.
Reads token from $IONOS_API_TOKEN (caller sources from ~/.ionos_token).
"""
from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
import time
import urllib.request
import urllib.error
from dataclasses import dataclass, asdict, field
from pathlib import Path

ENDPOINT = "https://openai.inference.de-txl.ionos.com/v1/chat/completions"
MODEL = "Qwen/Qwen3-Coder-Next"

REPO = Path(__file__).resolve().parents[2]
CMA = REPO / "experiments" / "2026-05-12-cross-model-authoring"
SPEC_CLASSIC = CMA / "rendered" / "ail.md"
SPEC_APPLY   = CMA / "rendered" / "ail-renamed.md"
SPEC_CALL    = CMA / "rendered" / "ail-call.md"
TASKS_DIR = CMA / "master" / "tasks"

# Reverse substitution per cohort so `ail check` (which only knows
# the classic tags) can parse outputs. Order: longer first.
REVERSE_APPLY = [
    ("tail-apply",      "tail-app"),
    ("term-constructor","term-ctor"),
    ("pat-constructor", "pat-ctor"),
    ("constructor",     "ctor"),
    ("apply",           "app"),
    ("lambda",          "lam"),
    ("typecon",         "con"),
]
REVERSE_CALL = [
    ("tail-call",       "tail-app"),
    ("term-constructor","term-ctor"),
    ("pat-constructor", "pat-ctor"),
    ("constructor",     "ctor"),
    ("call",            "app"),
    ("lambda",          "lam"),
    ("typecon",         "con"),
]

def reverse_substitute(src: str, pairs) -> str:
    out = src
    for new, old in pairs:
        out = re.sub(rf"\({re.escape(new)}(?=[\s)])", f"({old}", out)
    return out

# Tag sets used to count cohort-consistency.
TAGS_CLASSIC = ["tail-app","term-ctor","pat-ctor","ctor","app","lam","con"]
TAGS_APPLY   = ["tail-apply","term-constructor","pat-constructor","constructor","apply","lambda","typecon"]
TAGS_CALL    = ["tail-call","term-constructor","pat-constructor","constructor","call","lambda","typecon"]

def count_tags(src: str, tags: list[str]) -> dict[str,int]:
    counts: dict[str,int] = {}
    for t in tags:
        # match `(tag ` or `(tag)` only
        counts[t] = len(re.findall(rf"\({re.escape(t)}(?=[\s)])", src))
    return counts

@dataclass
class CallResult:
    cohort: str
    task_id: str
    trial: int
    prompt_tokens: int = 0
    completion_tokens: int = 0
    raw_response: str = ""
    extracted_code: str = ""
    own_tag_counts: dict[str,int] = field(default_factory=dict)
    foreign_tag_counts: dict[str,dict[str,int]] = field(default_factory=dict)
    ail_check_ok: bool = False
    ail_check_stderr: str = ""
    ail_build_ok: bool = False
    ail_run_ok: bool = False
    ail_run_stdout: str = ""
    stdout_matches_expected: bool = False
    notes: str = ""

def load_tasks() -> list[dict]:
    tasks = []
    for f in sorted(TASKS_DIR.glob("*.task.json")):
        tasks.append(json.loads(f.read_text()))
    return tasks

def extract_form_a(response: str) -> str:
    # Prefer fenced ```ail blocks; else fall back to the longest
    # parenthesised expression starting with (module …).
    m = re.findall(r"```(?:ail)?\s*\n(.*?)```", response, re.DOTALL)
    for block in m:
        if "(module" in block:
            return block.strip()
    # Fallback: from first '(module' to matching close paren.
    idx = response.find("(module")
    if idx < 0:
        return ""
    depth = 0
    out = []
    for ch in response[idx:]:
        out.append(ch)
        if ch == "(": depth += 1
        elif ch == ")":
            depth -= 1
            if depth == 0:
                return "".join(out)
    return "".join(out)

def call_ionos(token: str, system_prompt: str, user_prompt: str, *, max_retries=3) -> dict:
    body = json.dumps({
        "model": MODEL,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": user_prompt},
        ],
        "temperature": 0.2,
        "top_p": 0.95,
    }).encode("utf-8")
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    backoffs = [1, 4, 16]
    for attempt in range(max_retries):
        try:
            req = urllib.request.Request(ENDPOINT, data=body, headers=headers, method="POST")
            with urllib.request.urlopen(req, timeout=120) as r:
                return json.loads(r.read())
        except urllib.error.HTTPError as e:
            code = e.code
            if code in (429,) or 500 <= code < 600:
                if attempt < max_retries - 1:
                    time.sleep(backoffs[attempt])
                    continue
            raise
        except (urllib.error.URLError, TimeoutError) as e:
            if attempt < max_retries - 1:
                time.sleep(backoffs[attempt])
                continue
            raise
    raise RuntimeError("retries exhausted")

def run_pipeline(code: str, task_id: str, expected_stdout: str, work: Path) -> dict:
    src = work / f"{task_id}.ail"
    src.write_text(code)
    result = {"check_ok": False, "check_stderr": "", "build_ok": False,
              "run_ok": False, "run_stdout": "", "matches_expected": False}
    check = subprocess.run(["ail", "check", str(src)], capture_output=True, text=True, timeout=30)
    result["check_stderr"] = check.stderr + check.stdout if check.returncode != 0 else ""
    if check.returncode != 0:
        return result
    result["check_ok"] = True
    build = subprocess.run(["ail", "build", str(src), "-o", str(work / task_id)],
                           capture_output=True, text=True, timeout=60)
    if build.returncode != 0:
        result["check_stderr"] = build.stderr
        return result
    result["build_ok"] = True
    run = subprocess.run([str(work / task_id)], capture_output=True, text=True, timeout=10)
    if run.returncode == 0:
        result["run_ok"] = True
        result["run_stdout"] = run.stdout
        result["matches_expected"] = (run.stdout == expected_stdout)
    return result

def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--trials", type=int, default=2)
    ap.add_argument("--token-budget", type=int, default=300_000)
    ap.add_argument("--dry-run", action="store_true", help="don't call IONOS; use empty responses")
    args = ap.parse_args()

    token = os.environ.get("IONOS_API_TOKEN", "").strip()
    if not args.dry_run and not token:
        print("ERROR: IONOS_API_TOKEN not set", file=sys.stderr)
        return 2

    args.out.mkdir(parents=True, exist_ok=True)
    tasks = load_tasks()

    # (name, system_prompt, own_tags, reverse_pairs_or_None)
    cohorts = [
        ("classic", SPEC_CLASSIC.read_text(), TAGS_CLASSIC, None),
        ("apply",   SPEC_APPLY.read_text(),   TAGS_APPLY,   REVERSE_APPLY),
        ("call",    SPEC_CALL.read_text(),    TAGS_CALL,    REVERSE_CALL),
    ]
    all_tag_sets = {"classic": TAGS_CLASSIC, "apply": TAGS_APPLY, "call": TAGS_CALL}

    rows: list[CallResult] = []
    tokens_used = 0
    work = args.out / "work"
    work.mkdir(exist_ok=True)

    for cohort, system_prompt, own_tags, reverse_pairs in cohorts:
        for task in tasks:
            for trial in range(args.trials):
                if tokens_used >= args.token_budget:
                    print(f"BUDGET HIT at {tokens_used} tokens — stopping", file=sys.stderr)
                    break
                user_prompt = task["description"] + "\n\nReturn only the complete module as a single ```ail fenced block."
                row = CallResult(cohort=cohort, task_id=task["id"], trial=trial)
                if args.dry_run:
                    row.raw_response = "(dry-run)"
                else:
                    print(f"[{cohort}] {task['id']} trial {trial} …", file=sys.stderr)
                    resp = call_ionos(token, system_prompt, user_prompt)
                    row.raw_response = resp["choices"][0]["message"]["content"]
                    row.prompt_tokens = resp["usage"]["prompt_tokens"]
                    row.completion_tokens = resp["usage"]["completion_tokens"]
                    tokens_used += resp["usage"]["total_tokens"]
                code = extract_form_a(row.raw_response)
                row.extracted_code = code
                row.own_tag_counts = count_tags(code, own_tags)
                for other_name, other_tags in all_tag_sets.items():
                    if other_name != cohort:
                        row.foreign_tag_counts[other_name] = count_tags(code, other_tags)
                if code:
                    pipeline_input = reverse_substitute(code, reverse_pairs) if reverse_pairs else code
                    pipe = run_pipeline(pipeline_input, f"{cohort}_{task['id']}_t{trial}",
                                       task["expected_stdout"], work)
                    row.ail_check_ok = pipe["check_ok"]
                    row.ail_check_stderr = pipe["check_stderr"][:2000]
                    row.ail_build_ok = pipe["build_ok"]
                    row.ail_run_ok = pipe["run_ok"]
                    row.ail_run_stdout = pipe["run_stdout"]
                    row.stdout_matches_expected = pipe["matches_expected"]
                else:
                    row.notes = "no form-A extracted"
                rows.append(row)
                (args.out / f"{cohort}_{task['id']}_t{trial}.raw.txt").write_text(row.raw_response)

    out_file = args.out / "results.json"
    out_file.write_text(json.dumps([asdict(r) for r in rows], indent=2))
    print(f"\nwrote {out_file} — {len(rows)} rows, {tokens_used} tokens used")

    # Summary table.
    by_cohort: dict[str, dict] = {}
    for r in rows:
        c = by_cohort.setdefault(r.cohort, {"n":0, "extracted":0, "check":0, "build":0, "run":0, "correct":0,
                                            "own_tags":0, "foreign_tags":0})
        c["n"] += 1
        if r.extracted_code: c["extracted"] += 1
        if r.ail_check_ok:   c["check"] += 1
        if r.ail_build_ok:   c["build"] += 1
        if r.ail_run_ok:     c["run"] += 1
        if r.stdout_matches_expected: c["correct"] += 1
        c["own_tags"]     += sum(r.own_tag_counts.values())
        c["foreign_tags"] += sum(sum(v.values()) for v in r.foreign_tag_counts.values())

    print("\n=== summary ===")
    for cohort, c in by_cohort.items():
        print(f"{cohort:8s}  n={c['n']:2d}  extracted={c['extracted']:2d}  "
              f"check={c['check']:2d}  build={c['build']:2d}  run={c['run']:2d}  "
              f"correct={c['correct']:2d}  own-tags={c['own_tags']:3d}  foreign-tags={c['foreign_tags']:3d}")
    return 0

if __name__ == "__main__":
    sys.exit(main())