import openai

import torch

# api_key = "sk-804e05b936154e"
api_key = "asdfdfjadkf"
api_type = "open_ai"
api_version = "2024-11-01"
proxy = "None"
api_base = "http://{your llm deploy ip}:11434/v1"
LLM_MODEL = "qwen2:7b"


def chatCompletion(context):
    print("LLM - message:", flush=True)
    # 最后一个上下文的角色不能是system
    if len(context) > 0 and context[-1]['role'] == "system":
        context[-1]['role'] = "user"
    print(context, flush=True)

    openai.api_key = api_key
    openai.api_type = api_type
    openai.api_version = api_version
    openai.proxy = None if proxy == "None" else proxy
    openai_cli = openai.OpenAI(api_key=api_key, base_url=api_base)

    # 如果成功，返回响应文本、使用的总令牌数和一个表示成功的布尔值。如果出现异常，打印错误信息并重新抛出异常。
    try:
        response = openai_cli.chat.completions.create(
            model=LLM_MODEL,    # 指定使用的模型，比如常用的gpt-3.5-turbo等,这里是

            messages=context,   # 传入消息上下文
            max_tokens=2000,    # 限制生成回复的最大令牌（token）数，可根据需求调整
            temperature=0,  # 控制生成回复的创造性程度，取值范围0到1，接近0更保守，接近1更具创造性
            timeout=600     # 超时时间等参数
        )
        total_tokens = response.usage.total_tokens
        response_text = response.choices[0].message.content
        print("LLM - response_text:" + response_text, flush=True)
        return response_text, total_tokens, True
    except Exception as e:
        msg = "\nError: Failed to access GPT, please check whether your network can connect to GPT and terminal proxy is running properly.\n"
        print(f"\033[91m{msg} \033[0m")
        raise e


user_prompt = "如何做辣椒炒肉？"
final_context = [{"role": "system", "content": "You are a helpful assistant."},
                        {"role": "user", "content": user_prompt}]


message, total_tokens, success = chatCompletion(final_context)
print("total_tokens = %d" % total_tokens, flush=True)