#!/usr/bin/env python3
"""Замер скорости LLM через OpenAI (совместимый) API на двух фиксированных промптах.

Промпты:
  short — короткий запрос, но просит длинный ответ (~900 слов)
  long  — большой контекст ~100K токенов (статья, повторённая пронумерованными
          разделами; размер настраивается --long-context-tokens) + просьба
          длинного продолжения (~900 слов)

Метрики (стриминг, temperature=0):
  TTFT        — время до первого токена
  total       — общее время генерации
  tokens      — число сгенерированных токенов (из usage, либо оценка)
  tok/s (dec) — скорость декодирования: tokens / (total - TTFT)
  tok/s (e2e) — сквозная скорость: tokens / total

Запуск:
  export OPENAI_API_KEY=...
  python llm_benchmark.py --model gpt-4o-mini --runs 3

Примечания:
  - long-промпт ~100K токенов: модель должна поддерживать такой контекст.
  - Для точного подсчёта токенов: pip install tiktoken
    (иначе используется оценка ~4 символа/токен).
"""

from __future__ import annotations

import argparse
import math
import os
import statistics
import sys
import time
from dataclasses import dataclass

try:
    from openai import APIError, BadRequestError, OpenAI
except ImportError:
    sys.exit("Нужен пакет openai: pip install openai")


# ---------------------------------------------------------------------------
# Фиксированные промпты
# ---------------------------------------------------------------------------

SHORT_PROMPT = (
    "Write a very detailed, long story of about 900 words about a cat "
    "exploring an abandoned space station. Write it as a single flowing "
    "prose, no headings, no lists."
)

_LONG_DOC = """
ARTICLE: How to Benchmark Large Language Model Inference

Large language models are now used in everything from search engines to code assistants, and the practical question that follows is simple: how fast does a given model run? Speed is not a single number. A user waiting for the first sentence cares about time to first token, a user reading a long answer cares about tokens per second, and an operator running a batch job cares about total latency and cost. A good benchmark therefore measures several quantities at once and reports them separately instead of collapsing them into one ambiguous score.

Time to first token (TTFT) is the delay between sending a request and receiving the very first generated token. It is dominated by the time the model spends processing the input: the more context you send, the longer the prefill phase takes. TTFT is the metric that matters for interactive applications, because a human can start reading and feel that the system is responsive as soon as the first characters appear. If TTFT is several seconds, the interface feels broken even when the rest of the stream is fast.

Throughput, usually expressed in tokens per second, describes how quickly tokens arrive after the first one. This is the decoding phase, where the model generates the output one token at a time. Throughput determines how long a long answer takes to display. A model that produces 150 tokens per second finishes a one-thousand-token answer in roughly seven seconds of generation time; a model that produces 40 tokens per second needs twenty-five seconds for the same text. For a comfortable reading experience the floor is around twenty to thirty tokens per second.

The length of the prompt has a strong effect on both metrics. Prefill cost grows with the number of input tokens, so a short prompt with a long expected answer and a long prompt with a long expected answer exercise the system in different ways. A benchmark should include at least two shapes of workload: a short input that produces a long output, and a long input that also produces a long output. Comparing the two tells you whether the model scales gracefully with context or whether long documents make the whole pipeline collapse.

Sampling parameters change the generated text and can change the timing slightly, so a benchmark fixes them: temperature is set to zero, and the maximum number of output tokens is capped at a known value. With a fixed prompt and fixed parameters, repeated runs should produce nearly identical text and nearly identical timings. If they do not, the noise comes from the network, from shared hardware, or from server-side batching, and the benchmark should say so instead of hiding it.

Streaming is the right way to measure an LLM, because it is how real applications consume the model. With a streaming response you can timestamp each chunk: the first non-empty chunk gives TTFT, the last chunk gives total time, and the usage field at the end of the stream tells you exactly how many tokens were generated. A non-streaming call only gives you total latency, which mixes prefill and decoding into one number and makes it impossible to tell where the time went.

Token counting deserves attention because it is where naive benchmarks go wrong. The number of generated tokens is not the same as the number of words or characters, and estimating it by dividing the character count by four is only a rough approximation. Most APIs report the exact usage in the response; a serious benchmark should read that field and fall back to an estimate only when the field is missing. The estimate should be labeled as an estimate in the report, so nobody mistakes it for a measurement.

The first request after a cold start is usually slower than later ones: the server may load model weights, allocate memory, or warm up caches. A benchmark that includes the cold request in its average reports a number that will never be seen in production. The common practice is to run one warmup request and discard it, then measure a small number of real runs. Three to five runs is enough to get a stable picture without paying for a large number of tokens.

Run-to-run variance has several sources. The network adds jitter. The inference server batches requests from many users, so a run that lands next to a heavy batch is slower. The model itself, even at temperature zero, may produce slightly different text lengths because of sampling edges. A good report shows the minimum, the median, and the maximum of each metric across runs, so the reader can see both the typical case and the worst case.

Speed is only half of the story; the other half is cost. A fast model that costs ten times more per token may be the wrong choice for a high-volume feature, and a cheap model that is too slow may be unusable for an interactive one. A benchmark should record the price per one million input and output tokens and compute the estimated cost of each run, so that speed and money are compared on the same screen.

Prompt caching changes the picture for long inputs. If the server caches the prefix of a repeated prompt, the prefill phase becomes much cheaper and TTFT drops dramatically on the second and later runs. A benchmark that sends the same long prompt repeatedly is therefore measuring cached performance, not cold performance. If caching matters in your deployment, measure both: a first run with a fresh prompt and later runs with a cached one, and label the two numbers separately.

In summary, a useful LLM speed benchmark is boring on purpose: fixed prompts, fixed parameters, streaming responses, exact token counts, a warmup run, and a small number of measured runs per prompt shape. The report then shows, for each prompt, the time to first token, the total time, the number of output tokens, the tokens per second, and the estimated cost. Anyone reading the report should be able to answer, in one glance, whether the model is fast enough for the application in front of them.
"""

_LONG_TASK = (
    "\n\nTASK: Write a very detailed continuation and expansion of the "
    "article above, in the same style, of at least 900 words. Cover: "
    "hardware differences (GPU vs TPU), batching strategies, the effect "
    "of quantization on speed, autoregressive decoding vs speculative "
    "decoding, and practical tips for choosing a model. Continue with new "
    "material only; do not repeat the text above."
)

_LONG_INTRO = (
    "Below is a long technical article, split into numbered sections. "
    "Read all sections carefully before answering.\n"
)


def approx_tokens(text: str) -> int:
    """Число токенов: tiktoken (cl100k_base), если доступен, иначе ~4 симв./токен."""
    try:
        import tiktoken

        return len(tiktoken.get_encoding("cl100k_base").encode(text))
    except Exception:
        return max(1, round(len(text) / 4))


def build_long_prompt(target_tokens: int) -> str:
    """Промпт большого контекста: статья повторяется пронумерованными разделами,
    пока не наберётся ~target_tokens токенов, затем добавляется задача."""
    seed = _LONG_DOC.strip()
    seed_tokens = approx_tokens(seed)
    overhead = approx_tokens(_LONG_INTRO + _LONG_TASK)

    needed = max(seed_tokens, target_tokens - overhead)
    sections = math.ceil(needed / (seed_tokens + 6))  # +6 ≈ заголовок раздела

    parts = [_LONG_INTRO.rstrip("\n")]
    for i in range(1, sections + 1):
        parts.append(f"\n\n=== SECTION {i} ===\n\n{seed}")
    parts.append(_LONG_TASK)
    return "".join(parts)


def make_prompts(long_target_tokens: int) -> dict[str, str]:
    return {
        "short": SHORT_PROMPT,
        "long": build_long_prompt(long_target_tokens),
    }


# ---------------------------------------------------------------------------
# Замер одного запроса (стриминг)
# ---------------------------------------------------------------------------


@dataclass
class RunResult:
    ttft: float        # с, до первого токена
    total: float       # с, всего
    tokens: int        # сгенерированные токены (usage или оценка)
    estimated: bool    # True — токены оценены, а не измерены
    chars: int

    @property
    def decode_tps(self) -> float:
        """Скорость декодирования: токены/с после первого токена."""
        span = self.total - self.ttft
        return self.tokens / span if span > 0 else 0.0

    @property
    def e2e_tps(self) -> float:
        """Сквозная скорость: токены/с с момента отправки запроса."""
        return self.tokens / self.total if self.total > 0 else 0.0


def run_once(client: OpenAI, model: str, messages: list[dict[str, str]],
             max_tokens: int, temperature: float) -> RunResult:
    start = time.perf_counter()
    first_at: float | None = None
    parts: list[str] = []
    tokens: int | None = None

    def _create(with_usage: bool):
        kwargs = dict(
            model=model,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            stream=True,
        )
        if with_usage:
            kwargs["stream_options"] = {"include_usage": True}
        return client.chat.completions.create(**kwargs)

    try:
        stream = _create(with_usage=True)
    except BadRequestError:
        # Не все OpenAI-совместимые серверы понимают stream_options.
        stream = _create(with_usage=False)

    for chunk in stream:
        if chunk.choices:
            delta = chunk.choices[0].delta
            if delta is not None and delta.content:
                if first_at is None:
                    first_at = time.perf_counter()
                parts.append(delta.content)
        usage = getattr(chunk, "usage", None)
        if usage is not None and usage.completion_tokens:
            tokens = usage.completion_tokens

    total = time.perf_counter() - start
    text = "".join(parts)
    estimated = tokens is None
    if estimated:
        # Грубая оценка: в среднем ~4 символа на токен.
        tokens = max(1, round(len(text) / 4))

    return RunResult(
        ttft=(first_at - start) if first_at is not None else total,
        total=total,
        tokens=tokens,
        estimated=estimated,
        chars=len(text),
    )


# ---------------------------------------------------------------------------
# Замер промпта и отчёт
# ---------------------------------------------------------------------------


def bench(client: OpenAI, model: str, name: str, messages: list[dict[str, str]],
          runs: int, max_tokens: int, temperature: float) -> list[RunResult]:
    content = messages[-1]["content"]
    print(f"\n=== Промпт: {name} (вход ~{approx_tokens(content):,} токенов, "
          f"~{len(content):,} симв., max_tokens={max_tokens}) ===")

    print("  warmup (не учитывается)...")
    try:
        run_once(client, model, messages, max_tokens, temperature)
    except APIError as e:
        print(f"  warning: warmup не удался: {e}")

    results: list[RunResult] = []
    for i in range(runs):
        try:
            results.append(run_once(client, model, messages, max_tokens, temperature))
        except APIError as e:
            print(f"  warning: запуск {i + 1} не удался: {e}")
    return results


def report(name: str, results: list[RunResult]) -> None:
    if not results:
        print("  нет успешных запусков")
        return

    print(f"{'#':>3} | {'TTFT, с':>9} | {'всего, с':>9} | {'токены':>7} | "
          f"{'tok/s дек':>10} | {'tok/s e2e':>10}")
    for i, r in enumerate(results, 1):
        est = "*" if r.estimated else ""
        print(f"{i:>3} | {r.ttft:9.2f} | {r.total:9.2f} | {r.tokens:6d}{est} | "
              f"{r.decode_tps:10.1f} | {r.e2e_tps:10.1f}")

    mean_ttft = statistics.mean(r.ttft for r in results)
    mean_total = statistics.mean(r.total for r in results)
    mean_dec = statistics.mean(r.decode_tps for r in results)
    mean_e2e = statistics.mean(r.e2e_tps for r in results)
    med_ttft = statistics.median(r.ttft for r in results)
    med_dec = statistics.median(r.decode_tps for r in results)
    print(f"     среднее: TTFT {mean_ttft:.2f} с | всего {mean_total:.2f} с | "
          f"декод {mean_dec:.1f} tok/s | e2e {mean_e2e:.1f} tok/s")
    print(f"     медиана: TTFT {med_ttft:.2f} с | "
          f"декод {med_dec:.1f} tok/s")
    if any(r.estimated for r in results):
        print("     * число токенов оценено (usage API не отдаёт), а не измерено")


def main() -> None:
    parser = argparse.ArgumentParser(description="Замер скорости LLM через OpenAI API")
    parser.add_argument("--model",
                        default=os.environ.get("OPENAI_BENCH_MODEL", "gpt-4o-mini"),
                        help="имя модели (по умолчанию gpt-4o-mini)")
    parser.add_argument("--runs", type=int, default=3,
                        help="число замеров на промпт после warmup (по умолчанию 3)")
    parser.add_argument("--max-tokens", type=int, default=2048,
                        help="лимит на длину ответа (по умолчанию 2048)")
    parser.add_argument("--long-context-tokens", type=int, default=100_000,
                        help="целевой размер long-промпта в токенах (по умолчанию 100000)")
    parser.add_argument("--temperature", type=float, default=0.0,
                        help="температура (по умолчанию 0)")
    parser.add_argument("--base-url", default=None,
                        help="для OpenAI-совместимых API (vLLM, Ollama, ...)")
    parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"),
                        help="ключ API (по умолчанию из OPENAI_API_KEY)")
    args = parser.parse_args()

    if not args.api_key:
        if args.base_url:
            args.api_key = "EMPTY"  # локальные серверы часто не требуют ключ
        else:
            sys.exit("Не задан OPENAI_API_KEY (или --api-key)")

    client = OpenAI(api_key=args.api_key, base_url=args.base_url)

    print(f"Модель: {args.model} | runs={args.runs} | max_tokens={args.max_tokens} "
          f"| temperature={args.temperature}")
    if args.base_url:
        print(f"Base URL: {args.base_url}")

    prompts = make_prompts(args.long_context_tokens)
    long_in = approx_tokens(prompts["long"])
    print(f"Long-промпт: ~{long_in:,} токенов на запрос; на весь прогон "
          f"(warmup + {args.runs} runs) уйдёт ~{long_in * (args.runs + 1):,} "
          f"входных токенов — учтите стоимость и лимит контекста модели.")

    results_by_name: dict[str, list[RunResult]] = {}
    for name in ("short", "long"):
        messages = [{"role": "user", "content": prompts[name]}]
        results_by_name[name] = bench(
            client, args.model, name, messages,
            args.runs, args.max_tokens, args.temperature,
        )
        report(name, results_by_name[name])

    short, long_ = results_by_name["short"], results_by_name["long"]
    if short and long_:
        s_ttft = statistics.mean(r.ttft for r in short)
        l_ttft = statistics.mean(r.ttft for r in long_)
        s_dec = statistics.mean(r.decode_tps for r in short)
        l_dec = statistics.mean(r.decode_tps for r in long_)
        print(f"\nСравнение: TTFT short {s_ttft:.2f} с vs long {l_ttft:.2f} с "
              f"(разница — цена префилла ~{long_in:,} токенов контекста; после "
              f"warmup сервер может кэшировать префикс, тогда TTFT повторных "
              f"запусков занижен); декодирование: {s_dec:.1f} vs {l_dec:.1f} tok/s")

    if not all(results_by_name.values()):
        sys.exit(1)


if __name__ == "__main__":
    main()
