Measure what happens when requests overlap

import argparse
from concurrent.futures import ThreadPoolExecutor
from time import perf_counter
import statistics

from compatible_chat import local_base, request

PROMPT = (
    "Write a plain-language paragraph with six practical suggestions "
    "for planning a community library workshop."
)


def run_request(base, model, prompt):
    started = perf_counter()
    reply = request(base, "/chat/completions", {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0,
        "max_tokens": 96,
        "stream": False,
    })
    duration = perf_counter() - started
    content = reply["choices"][0]["message"].get("content")
    if not isinstance(content, str) or not content.strip():
        raise ValueError("A request returned no visible text.")
    return duration, reply.get("usage", {}).get("completion_tokens")


def run_batch(base, model, clients, count):
    prompts = [f"{PROMPT} Start with the room booking." if index % 2 else PROMPT
               for index in range(count)]
    start = perf_counter()
    with ThreadPoolExecutor(max_workers=clients) as pool:
        results = list(pool.map(lambda prompt: run_request(base, model, prompt), prompts))
    wall = perf_counter() - start
    print(f"{clients} client(s): {count} replies in {wall:.2f} s")
    print(f"  Median reply time: {statistics.median(row[0] for row in results):.2f} s")
    print(f"  Throughput: {count / wall:.2f} replies/s")
    tokens = [row[1] for row in results]
    if all(isinstance(value, int) for value in tokens):
        print(f"  Completion tokens: {sum(tokens)}")


parser = argparse.ArgumentParser(description="Compare one and two concurrent local clients.")
parser.add_argument("--base", type=local_base, default="http://127.0.0.1:11434/v1")
parser.add_argument("--model", required=True)
parser.add_argument("--count", type=int, default=4, choices=(4, 6, 8))
parser.add_argument("--parallel-first", action="store_true")
args = parser.parse_args()

try:
    run_request(args.base, args.model, "Name one library service.")
    for clients in ((2, 1) if args.parallel_first else (1, 2)):
        run_batch(args.base, args.model, clients, args.count)
except (OSError, KeyError, IndexError, TypeError, ValueError) as error:
    raise SystemExit(f"Could not complete the measurement: {error}") from None

import argparse
import json
import os
from urllib.error import HTTPError
from urllib.parse import urlsplit
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener


class NoRedirects(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def local_base(value):
    url = urlsplit(value)
    if (
        url.scheme != "http"
        or url.hostname not in {"localhost", "127.0.0.1", "::1"}
        or url.username is not None
        or url.password is not None
        or url.path.rstrip("/") != "/v1"
        or url.query
        or url.fragment
    ):
        raise argparse.ArgumentTypeError("Use a loopback HTTP address ending in /v1.")
    return value.rstrip("/")


def request(base, endpoint, payload=None, *, timeout=120):
    headers = {"Content-Type": "application/json"}
    key = os.environ.get("LOCAL_MODEL_API_KEY")
    if key:
        headers["Authorization"] = f"Bearer {key}"
    data = json.dumps(payload).encode("utf-8") if payload is not None else None
    message = Request(base + endpoint, data=data, headers=headers)
    opener = build_opener(ProxyHandler({}), NoRedirects())
    with opener.open(message, timeout=timeout) as response:
        return json.load(response)


def main():
    parser = argparse.ArgumentParser(description="Talk to a local Chat Completions endpoint.")
    parser.add_argument("--base", type=local_base, default="http://127.0.0.1:11434/v1")
    parser.add_argument("--list", action="store_true", help="Print available model IDs")
    parser.add_argument("--model", help="Exact model ID from --list")
    parser.add_argument("prompt", nargs="?")
    args = parser.parse_args()
    if not args.list and (not args.model or not args.prompt):
        parser.error("Use --list, or provide --model and a prompt.")

    try:
        if args.list:
            for model in request(args.base, "/models")["data"]:
                print(model["id"])
            return

        result = request(args.base, "/chat/completions", {
            "model": args.model,
            "messages": [{"role": "user", "content": args.prompt}],
            "temperature": 0,
            "max_tokens": 96,
            "stream": False,
        })
        text = result["choices"][0]["message"].get("content")
        if not isinstance(text, str) or not text.strip():
            raise ValueError("The server returned no visible reply. Check the model and output limit.")
        print(text)
    except HTTPError as error:
        raise SystemExit(f"Server returned HTTP {error.code}. Check the model ID, port, and API key.") from None
    except (OSError, KeyError, IndexError, TypeError, ValueError) as error:
        raise SystemExit(f"Could not read a local reply: {error}") from None


if __name__ == "__main__":
    main()