from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate

from langchain_openai import ChatOpenAI

from env_util import DASHSCOPE_API_KEY, DASHSCOPE_BASE_URL

# llm对象就是调用大模型的对象
llm = ChatOpenAI(
    model = "qwen-plus",
    base_url=DASHSCOPE_BASE_URL,
    api_key=DASHSCOPE_API_KEY,
    temperature=0.8,
);

# 提示词
examples = [
    {
        "question": "罗纳尔多·路易斯·纳萨里奥·德·利马 世界杯进了多少球？",
        "answer": """
            中间问题：罗纳尔多·路易斯·纳萨里奥·德·利马 参加了哪几届世界杯？
            中间答案：…………
            中间问题：罗纳尔多·路易斯·纳萨里奥·德·利马 分别在这几届进了多少球？
            中间答案：…………
            最终答案：罗纳尔多·路易斯·纳萨里奥·德·利马 世界杯进了xx球。
        """
    },
    {
        "question": "周星驰参演了多少部电影？",
        "answer": """
            中间问题：周星驰哪几年参演了电影？
            中间答案：…………
            中间问题：周星驰分别在这几年参演了多少电影？
            中间答案：…………
            最终答案：周星驰参演了xx部电影。
        """
    }
]

base_template = PromptTemplate.from_template("问题：{question}\n{answer}")

example_template = FewShotPromptTemplate(
    example_prompt = base_template,  # 前面的模板格式
    examples = examples,       # ICL的示例
    suffix = "问题：{input}",
    input_variables=["input"]
)

chain = example_template | llm

input = {"input": "梅西世界杯进了多少球？"}

resp = chain.invoke(input)

print(resp)