lmstudiofixed/tools/lmstudio/probes/llama_cpp_openai_probe.py

#!/usr/bin/env python3
import argparse
import io
import json
import math
import struct
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path


CHAT_TESTS = [
    ("baseline", {}),
    ("enable_thinking_false", {"enable_thinking": False}),
    ("reasoning_off", {"reasoning": "off"}),
]


def post_json(url: str, payload: dict, timeout: int) -> tuple[int, str]:
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read().decode("utf-8", errors="replace")


def post_multipart(url: str, fields: dict[str, str], file_name: str, file_bytes: bytes, timeout: int) -> tuple[int, str]:
    boundary = "----CodexBoundary20260609"
    body = io.BytesIO()

    for key, value in fields.items():
        body.write(f"--{boundary}\r\n".encode())
        body.write(f'Content-Disposition: form-data; name="{key}"\r\n\r\n'.encode())
        body.write(f"{value}\r\n".encode())

    body.write(f"--{boundary}\r\n".encode())
    body.write(
        f'Content-Disposition: form-data; name="file"; filename="{file_name}"\r\n'.encode()
    )
    body.write(b"Content-Type: audio/wav\r\n\r\n")
    body.write(file_bytes)
    body.write(b"\r\n")
    body.write(f"--{boundary}--\r\n".encode())

    req = urllib.request.Request(
        url,
        data=body.getvalue(),
        headers={"Content-Type": f"multipart/form-data; boundary={boundary}"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read().decode("utf-8", errors="replace")
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read().decode("utf-8", errors="replace")


def generate_silence_wav(sample_rate: int = 16000, duration_seconds: float = 1.0) -> bytes:
    total_samples = int(sample_rate * duration_seconds)
    pcm = b"".join(struct.pack("<h", 0) for _ in range(total_samples))
    byte_rate = sample_rate * 2
    block_align = 2
    bits_per_sample = 16
    data_size = len(pcm)
    riff_size = 36 + data_size

    header = struct.pack(
        "<4sI4s4sIHHIIHH4sI",
        b"RIFF",
        riff_size,
        b"WAVE",
        b"fmt ",
        16,
        1,
        1,
        sample_rate,
        byte_rate,
        block_align,
        bits_per_sample,
        b"data",
        data_size,
    )
    return header + pcm


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=18003)
    parser.add_argument("--model", default="Qwen3-ASR-1.7B-Q8_0.gguf")
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--output-root", default="artifacts/llama-cpp-openai-probe")
    args = parser.parse_args()

    run_dir = Path(args.output_root) / datetime.now().strftime("%Y%m%d-%H%M%S")
    run_dir.mkdir(parents=True, exist_ok=True)

    chat_url = f"http://127.0.0.1:{args.port}/v1/chat/completions"
    audio_url = f"http://127.0.0.1:{args.port}/v1/audio/transcriptions"

    chat_results = []
    for index, (name, extra) in enumerate(CHAT_TESTS, start=1):
        payload = {
            "model": args.model,
            "messages": [{"role": "user", "content": "Hello"}],
            "temperature": 0.2,
            "max_tokens": 64,
            "stream": False,
        }
        payload.update(extra)
        status, raw = post_json(chat_url, payload, timeout=args.timeout)
        (run_dir / f"chat_{index:02d}_{name}.request.json").write_text(
            json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
        (run_dir / f"chat_{index:02d}_{name}.response.raw.json").write_text(raw, encoding="utf-8")
        chat_results.append({"name": name, "status": status, "raw": raw})

    wav_bytes = generate_silence_wav()
    status, raw = post_multipart(
        audio_url,
        fields={"model": args.model},
        file_name="silence.wav",
        file_bytes=wav_bytes,
        timeout=args.timeout,
    )
    (run_dir / "audio_silence.response.raw.json").write_text(raw, encoding="utf-8")
    (run_dir / "audio_silence.summary.json").write_text(
        json.dumps({"status": status, "raw": raw}, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )

    (run_dir / "all_summaries.json").write_text(
        json.dumps(
            {
                "chat_results": chat_results,
                "audio_result": {"status": status, "raw": raw},
            },
            ensure_ascii=False,
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    print(run_dir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())