from langchain_core.messages import HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, \
    FewShotChatMessagePromptTemplate

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 = [
    {"input": "2 😫 3","output": "6"},
    {"input": "2 😫 4","output": "8"},
];
# 示例格式（给大模型学习的）
example_prompt = ChatPromptTemplate.from_messages(
    [
        ('human',"{input}"),
        ('ai',"{output}")
    ]
);
# 组装ICL模板
icl_template = FewShotChatMessagePromptTemplate(
    examples=examples,
    example_prompt=example_prompt,
);

# 聊天模板，组合上ICL
prompt_template = ChatPromptTemplate.from_messages([
    ("system","你是一个智能助手"),
    icl_template,
    MessagesPlaceholder("ques")
]);

# 组装执行
chain = prompt_template | llm

resp = chain.invoke(
    # 课程里说可以不加[]，但是测试得出，[]必须要有。。。。
    {"ques": [HumanMessage(content = "2 😫 67")]}
);

print(resp)