from openai import OpenAI
import sys
from prompt_toolkit import PromptSession
from prompt_toolkit.styles import Style
import time
import tiktoken
import shutil

# 初始化 OpenAI 客户端，base_url,api_key,model请按实际情况填写。
client = OpenAI(
    base_url=<填写魔搭侧的服务地址>,
    api_key=<填写魔搭侧的服务api_key>
)

model = <填写模型路径，如unsloth/Qwen3-8B-GGUF>

def num_tokens_from_string(string: str) -> int:
    """计算文本中的 tokens 数量"""
    encoding = tiktoken.get_encoding("cl100k_base")
    num_tokens = len(encoding.encode(string))
    return num_tokens

def get_terminal_size():
    """获取终端大小"""
    return shutil.get_terminal_size()

def clear_current_line():
    """清除当前行"""
    terminal_width = get_terminal_size().columns
    print('\r' + ' ' * terminal_width + '\r', end='', flush=True)

def chat_with_model():
    # 创建 prompt session 和自定义样式
    session = PromptSession()
    style = Style.from_dict({
        'prompt': '#00aa00 bold',
    })
    
    messages = [
        {
            'role': 'system',
            'content': '你是一个有帮助的助手。'
        }
    ]
    
    print("欢迎使用AI助手！输入'退出'或'quit'可以结束对话。")
    
    while True:
        try:
            # 使用 prompt_toolkit 处理输入
            user_input = session.prompt('\n你: ', style=style)
            
            if user_input.lower() in ['退出', 'quit']:
                print("再见！")
                break
                
            messages.append({
                'role': 'user',
                'content': user_input
            })
            
            print("\nAI助手: ", end='', flush=True)
            try:
                start_time = time.time()
                response = client.chat.completions.create(
                    model=model,
                    messages=messages,
                    stream=True
                )
                
                full_response = ""
                for chunk in response:
                    if chunk.choices[0].delta.content:
                        content = chunk.choices[0].delta.content
                        print(content, end='', flush=True)
                        full_response += content
                
                end_time = time.time()
                elapsed_time = end_time - start_time
                
                # 计算 tokens
                response_tokens = num_tokens_from_string(full_response)
                tokens_per_second = response_tokens / elapsed_time
                
                # 确保统计信息显示在新行，并在显示后额外添加换行
                print("\n")  # 先确保有一个空行
                stats = f"[统计] 响应tokens: {response_tokens}, 用时: {elapsed_time:.2f}秒, 速率: {tokens_per_second:.2f} tokens/s"
                print(stats)
                print()  # 添加额外的换行确保滚动
                sys.stdout.flush()  # 强制刷新输出缓冲区
                
                messages.append({
                    'role': 'assistant',
                    'content': full_response
                })
                
            except Exception as api_error:
                print(f"\n请求API时发生错误: {str(api_error)}")
                print()  # 额外的换行
                continue
            
        except KeyboardInterrupt:
            print("\n\n检测到中断，正在退出...")
            print()  # 额外的换行
            break
        except Exception as e:
            print(f"\n发生错误: {str(e)}")
            print()  # 额外的换行
            continue

if __name__ == "__main__":
    try:
        chat_with_model()
    except Exception as e:
        print(f"程序发生错误: {str(e)}")
        print()  # 额外的换行
        sys.exit(1)