import os
from openai import OpenAI
from langchain.embeddings.base import Embeddings
from dotenv import load_dotenv

# 加载 .env 文件中的环境变量, 隐藏 API Keys
load_dotenv()


# ============================================================
# 模型1: 嵌入模型, 采用 OpenAI text-embedding-3-small
# ============================================================

class OpenAIEmbeddings(Embeddings):
    """基于 OpenAI Embedding API 的自定义嵌入类"""

    def __init__(self):
        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    def embed_documents(self, texts):
        embeddings = []
        for text in texts:
            response = self.client.embeddings.create(
                model="text-embedding-3-small",
                input=[text],
            )
            embeddings.append(response.data[0].embedding)
        return embeddings

    def embed_query(self, text):
        # 查询文档
        return self.embed_documents([text])[0]


# ============================================================
# 模型2: 本地大语言模型, 通过 vLLM 服务 (OpenAI 兼容接口)
# 启动命令: vllm serve ./Qwen3-Next-80B-A3B-Thinking-AWQ-4bit --dtype auto --trust-remote-code --max-model-len 4096 --port 8000
# ============================================================

VLLM_MODEL_NAME = "./Qwen3-Next-80B-A3B-Thinking-AWQ-4bit"
VLLM_BASE_URL = "http://localhost:8000/v1"


def create_local_llm_client():
    """创建本地 vLLM 客户端 (OpenAI 兼容接口)"""
    client = OpenAI(api_key="none", base_url=VLLM_BASE_URL)
    return client


def generate_local_answer(client, question):
    """调用本地 vLLM 生成回答"""
    response = client.chat.completions.create(
        model=VLLM_MODEL_NAME,
        messages=[
            {"role": "system", "content": "你是一个能力非常强大的助手."},
            {"role": "user", "content": question}
        ],
        max_tokens=2048,
        stream=False
    )
    return response.choices[0].message.content


# ============================================================
# 模型3: 远程大语言模型, 采用 ChatGPT (OpenAI API)
# ============================================================

def create_chatgpt_client():
    """创建 OpenAI ChatGPT 客户端"""
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    return client


def generate_chatgpt_answer(client, question):
    """调用 ChatGPT 生成回答"""
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "你是一个能力非常强大的助手."},
            {"role": "user", "content": question}
        ],
        stream=False
    )
    return response.choices[0].message.content


# ============================================================
# 测试入口
# ============================================================

if __name__ == "__main__":
    # 测试本地 vLLM 模型
    client = create_local_llm_client()
    output = generate_local_answer(client, "你好啊,千与千寻")
    print('-' * 50)
    print(output)

    # 测试远程 ChatGPT
    #client = create_chatgpt_client()
    #output = generate_chatgpt_answer(client, "你好啊,千与千寻")
    #print('-' * 50)
    #print(output)