#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""专题五 · 步骤 1：验证 vLLM 三层（列模型 / 单轮对话 / 工具调用）。"""
from __future__ import annotations

import json
import sys

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


def curl_json(d: Dify, cmd: str, label: str) -> None:
    out, err = d.sh(cmd, t=120)
    print("\n=== %s ===" % label)
    print((out or err)[:1500])


def main() -> int:
    d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD)
    try:
        d.login()
        auth = "-H 'Authorization: Bearer %s'" % C.VLLM_KEY
        curl_json(
            d,
            "curl -s %s/v1/models %s | head -c 800" % (C.VLLM_URL, auth),
            "3.1 GET /v1/models",
        )
        chat_body = {
            "model": C.LLM_MODEL,
            "messages": [
                {"role": "system", "content": "你是一个助手"},
                {"role": "user", "content": "你好"},
            ],
            "max_tokens": 128,
        }
        remote = "/tmp/vllm_chat_smoke.json"
        with d.sftp.open(remote, "w") as f:
            f.write(json.dumps(chat_body, ensure_ascii=False))
        curl_json(
            d,
            "curl -s -X POST %s/v1/chat/completions -H 'Content-Type: application/json' %s -d @%s"
            % (C.VLLM_URL, auth, remote),
            "3.2 POST /v1/chat/completions",
        )
        tool_body = {
            "model": C.LLM_MODEL,
            "messages": [{"role": "user", "content": "上海今天天气怎么样"}],
            "tools": [{
                "type": "function",
                "function": {
                    "name": "get_weather",
                    "description": "查询天气",
                    "parameters": {
                        "type": "object",
                        "properties": {"city": {"type": "string"}},
                        "required": ["city"],
                    },
                },
            }],
            "tool_choice": "auto",
        }
        remote2 = "/tmp/vllm_tool_smoke.json"
        with d.sftp.open(remote2, "w") as f:
            f.write(json.dumps(tool_body, ensure_ascii=False))
        curl_json(
            d,
            "curl -s -X POST %s/v1/chat/completions -H 'Content-Type: application/json' %s -d @%s"
            % (C.VLLM_URL, auth, remote2),
            "3.3 POST /v1/chat/completions（带 tools）",
        )
        return 0
    finally:
        d.close()


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