#!/usr/bin/env python3
"""
交互式对话脚本，用于与vLLM模型进行对话并计算token生成速度

使用方法：
1. 进入Docker容器：docker exec -it rocm-vllm-openai bash
2. 运行：python chat_with_speed.py
3. 输入提示词与模型对话
4. 输入 'exit' 退出
"""

import json
import time
import httpx

# 模型服务地址
API_URL = "http://localhost:8001/v1/chat/completions"
# API密钥
API_KEY = "sk-szcjw"
# 模型名称
MODEL_NAME = "Qwen_local_model"

def chat_with_model():
    """交互式对话函数"""
    print("=== vLLM 交互式对话工具 ===")
    print("输入提示词与模型对话，输入 'exit' 退出")
    print("=" * 50)
    
    # 对话历史
    messages = []
    
    while True:
        # 获取用户输入
        user_input = input("用户: ").strip()
        
        if user_input.lower() == "exit":
            print("退出对话...")
            break
        
        if not user_input:
            continue
        
        # 添加用户消息到对话历史
        messages.append({"role": "user", "content": user_input})
        
        # 准备请求数据
        payload = {
            "model": MODEL_NAME,
            "messages": messages,
            "max_tokens": 1000,
            "temperature": 0.7,
            "top_p": 0.8,
            "top_k": 20
        }
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {API_KEY}"
        }
        
        print("模型: ", end="", flush=True)
        
        # 记录开始时间
        start_time = time.time()
        
        try:
            # 发送请求
            response = httpx.post(API_URL, json=payload, headers=headers, timeout=300.0)
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            
            # 获取模型回复
            assistant_message = result["choices"][0]["message"]["content"]
            print(assistant_message)
            
            # 添加模型回复到对话历史
            messages.append({"role": "assistant", "content": assistant_message})
            
            # 计算token速度
            usage = result.get("usage", {})
            completion_tokens = usage.get("completion_tokens", 0)
            end_time = time.time()
            elapsed_time = end_time - start_time
            
            if completion_tokens > 0 and elapsed_time > 0:
                tokens_per_second = completion_tokens / elapsed_time
                print(f"\n[速度统计] 生成 {completion_tokens} tokens，用时 {elapsed_time:.2f} 秒，速度: {tokens_per_second:.2f} tokens/s")
            else:
                print("\n[速度统计] 无法计算速度")
                
        except Exception as e:
            print(f"\n错误: {e}")
            
        print("=" * 50)

if __name__ == "__main__":
    chat_with_model()