import os

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI

load_dotenv()

# 1.准备大模型对象（客户端）
client = ChatOpenAI(base_url=os.getenv("BASE_URL"),
                    api_key=os.getenv("DEEPSEEK_API_KEY"),
                    model=os.getenv("MODEL_NAME"),
                    temperature=0)

# 2.准备提示词
class_examples = {
    '新闻报道': '今日，股市经历了一轮震荡，受到宏观经济数据和全球贸易紧张局势的影响。投资者密切关注美联储可能的政策调整，以适应市场的不确定性。',
    '财务报告': '本公司年度财务报告显示，去年公司实现了稳步增长的盈利，同时资产负债表呈现强劲的状况。经济环境的稳定和管理层的有效战略执行为公司的健康发展奠定了基础。',
    '公司公告': '本公司高兴地宣布成功完成最新一轮并购交易，收购了一家在人工智能领域领先的公司。这一战略举措将有助于扩大我们的业务领域，提高市场竞争力。',
    '分析师报告': '最新的行业分析报告指出，科技公司的创新将成为未来增长的主要推动力。云计算、人工智能和数字化转型被认为是引领行业发展的关键因素，投资者应关注这些趋势。'}


# 把上述的提示词文本，添加到提示词的案例（few-shot）中
def build_prompt():
    # 获取所有的类别
    classes = class_examples.keys()
    # 初始化提示
    init_prompt = f"你是一个文本分类专家，请将文本划分为{classes}类别中的某个类别。示例如下：\n"
    # 遍历示例
    for item, content in class_examples.items():
        init_prompt += f"{content}是{classes}中的{item}类别。\n"
    # 输出提示词
    return init_prompt


sentences = [
    "今日，央行发布公告宣布降低利率，以刺激经济增长。这一降息举措将影响贷款利率，并在未来几个季度内对金融市场产生影响。",
    "ABC公司今日发布公告称，已成功完成对XYZ公司股权的收购交易。本次交易是ABC公司在扩大业务范围、加强市场竞争力方面的重要举措。据悉，此次收购将进一步巩固ABC公司在行业中的地位，并为未来业务发展提供更广阔的发展空间。详情请见公司官方网站公告栏。",
    "公司资产负债表显示，公司偿债能力强劲，现金流充足，为未来投资和扩张提供了坚实的财务基础。",
    "最新的分析报告指出，可再生能源行业预计将在未来几年经历持续增长，投资者应该关注这一领域的投资机会。",
]


# 3.把提示词送给大模型
def call_llm(prompt):
    # 拼接提示词
    prompt += "请对如下文本进行分类，输出最终类别即可。\n"
    # 得到文本
    for sentence in sentences:
        # 调用大模型
        result = client.invoke(prompt + sentence)

        # 4.得到结果
        print("result-->", result.content)


if __name__ == '__main__':
    init_prompt = build_prompt()
    # print("init_prompt-->", init_prompt)
    call_llm(init_prompt)