from dotenv import load_dotenv
import os
from langchain_openai import ChatOpenAI

# 加载配置
load_dotenv()

# 1.创建大模型客户端
client = ChatOpenAI(model=os.getenv("MODEL_NAME"),
                    base_url=os.getenv("BASE_URL"),
                    api_key=os.getenv("DEEPSEEK_API_KEY"),
                    temperature=0)

schema = {"金融": ["日期", "股票名称", "开盘价", "收盘价", "成交量"]}

examples = {'金融': [{"content": "2023-01-10，股市震荡。股票古哥-D[EOOE]美股今日开盘价100美元，一度飙升至105美元，随后回落至98美元，最终以102美元收盘，成交量达到520000。"},
                    {"answer": {"日期": "2023-01-10", '股票名称': "古哥-D[EOOE]美股", '开盘价': "100美元", '收盘价': "102美元", '成交量': "520000"}}]}

sentences = ['2023-02-15，寓意吉祥的节日，股票佰笃[BD]美股开盘价10美元，虽然经历了波动，但最终以13美元收盘，成交量微幅增加至460,000，投资者情绪较为平稳。',
'2023-04-05，市场迎来轻松氛围，股票盘古(0021)开盘价23元，尽管经历了波动，但最终以26美元收盘，成交量缩小至310,000，投资者保持观望态度。',]

def build_prompt():
    init_prompt = f"你是一个文本信息抽取专家。参考格式：{schema}，完成对输入的文本进行信息抽取任务。比如："
    for example in examples.items():
        init_prompt += str(example)
    # print("init_prompt-->", init_prompt)
    return init_prompt


def call_llm(prompt):
    prompt += '请对如下文本进行信息抽取，只需要输出最终的结果即可，不需要输出其他内容。\n'
    for sentence in sentences:
        # print("sentence-->", sentence)
        result= client.invoke(prompt + sentence)
        print(result.content)


if __name__ == '__main__':
    prompt = build_prompt()
    call_llm(prompt)