#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题十四 · 步骤 3：分层探针（Xinference / Rerank / vLLM / 最短对话）。"""
from __future__ import annotations

import json
import sys

import _bootstrap  # noqa: F401
import _config as C
from dify_api import Dify


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        probes = [
            ("1) Xinference /v1/models", (
                "curl -s -m 15 -H 'Authorization: Bearer %s' %s/v1/models"
                % (C.XF_KEY, C.XF_URL)
            )),
            ("2) Xinference /v1/rerank", (
                "curl -s -m 30 -H 'Authorization: Bearer %s' -H 'Content-Type: application/json' "
                "-d '{\"model\":\"%s\",\"query\":\"%s\",\"documents\":[\"管理人员应通过钉钉打卡。\",\"班车时间以通知为准。\"]}' "
                "%s/v1/rerank"
                % (C.XF_KEY, C.RERANK_MODEL, C.PROBE_QUERY_KB, C.XF_URL)
            )),
            ("3) vLLM /v1/models", (
                "curl -s -m 10 -H 'Authorization: Bearer %s' %s/models"
                % (C.VLLM_KEY, C.VLLM_BASE)
            )),
            ("4) vLLM 最短对话", (
                "curl -s -m 60 -X POST %s/chat/completions "
                "-H 'Authorization: Bearer %s' -H 'Content-Type: application/json' "
                "-d '{\"model\":\"%s\",\"messages\":[{\"role\":\"user\",\"content\":\"只回：pong\"}],\"max_tokens\":16}'"
                % (C.VLLM_BASE, C.VLLM_KEY, C.QWEN_MODEL)
            )),
        ]
        for title, cmd in probes:
            print("\n=== %s ===" % title)
            out, err = d.sh(cmd)
            text = out or err
            print(text[:1500] if len(text) > 1500 else text)
        return 0
    finally:
        d.close()


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