scripts/smoke_llm.py

#!/usr/bin/env python3
"""Smoke test for the LLM endpoint, before wiring up the whole app.

    python scripts/smoke_llm.py                    # stream một câu chào
    python scripts/smoke_llm.py --lesson "Pytago"  # sinh giáo án thật, in ra JSON

Streaming here is only for eyeballing latency and checking the key works. The
API itself does not stream: a lesson has to be complete and validated before the
board can render step one, so there is nothing useful to show halfway through.
"""

import argparse
import json
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "backend"))

from openai import OpenAI  # noqa: E402

from app.config import settings  # noqa: E402

_COLOR = sys.stdout.isatty() and os.getenv("NO_COLOR") is None
DIM = "\033[90m" if _COLOR else ""
RESET = "\033[0m" if _COLOR else ""


def stream(prompt: str) -> None:
    client = OpenAI(base_url=settings.llm_base_url, api_key=settings.llm_api_key)
    kwargs = {
        "model": settings.model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": settings.temperature,
        "top_p": settings.top_p,
        "max_tokens": settings.max_tokens,
        "stream": True,
    }
    if settings.seed is not None:
        kwargs["seed"] = settings.seed

    reasoning_open = False
    for chunk in client.chat.completions.create(**kwargs):
        if not getattr(chunk, "choices", None):
            continue
        delta = getattr(chunk.choices[0], "delta", None)
        if delta is None:
            continue

        # GLM và các model reasoning khác tách phần suy nghĩ ra kênh riêng.
        thought = getattr(delta, "reasoning_content", None)
        if thought:
            if not reasoning_open:
                print(DIM, end="")
                reasoning_open = True
            print(thought, end="", flush=True)

        if delta.content:
            if reasoning_open:
                print(RESET, end="")
                reasoning_open = False
            print(delta.content, end="", flush=True)

    print(RESET if reasoning_open else "")


def lesson(question: str) -> None:
    from app.services.llm import lesson_service

    result = lesson_service.generate(question)
    print(f"chế độ structured output: {lesson_service._mode}\n")
    print(json.dumps(result.model_dump(), ensure_ascii=False, indent=2))


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--lesson", metavar="CÂU_HỎI", help="sinh giáo án thay vì stream")
    parser.add_argument("--prompt", default="Chào bạn, giới thiệu ngắn về bản thân bằng tiếng Việt.")
    args = parser.parse_args()

    if not settings.llm_ready:
        print("Chưa có LLM_API_KEY. Đặt trong .env rồi thử lại.", file=sys.stderr)
        return 1

    print(f"endpoint: {settings.llm_base_url}\nmodel:    {settings.model}\n")
    if args.lesson:
        lesson(args.lesson)
    else:
        stream(args.prompt)
    return 0


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