import logging

from langchain_core.prompts import FewShotChatMessagePromptTemplate, ChatPromptTemplate
from langchain_openai import ChatOpenAI
import os
import dotenv

dotenv.load_dotenv()

## 设置环境变量
os.environ['OPENAI_API_KEY'] = os.getenv("SILICONFLOW_API_KEY")
os.environ['OPENAI_BASE_URL'] = os.getenv("SILICONFLOW_BASE_URL")

# 默认的 'model_name': 'deepseek-ai/DeepSeek-V3.1',
llm = ChatOpenAI(model="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B")

logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)

examples = [
    {"input": "1 || 1", "output": "2"},
    {"input": "1 || 2", "output": "3"},
    {"input": "1 || 3", "output": "4"}
]

example_prompt = ChatPromptTemplate.from_messages(
    [
        ("human", "{input}"),
        ("ai", "{output}"),
    ]
)

few_show_prompt = FewShotChatMessagePromptTemplate(
    examples=examples,
    example_prompt=example_prompt
)


final_prompt = ChatPromptTemplate.from_messages(
    [
        ("system","你是一个数学天才"),
        few_show_prompt,
        ("human", "{input}")
    ]
)
# question = final_prompt.invoke(input = {"input":"1 || 10"})
# # llm  : 1||10  ?
# response = llm.invoke(question)
# print(response)
# 链式调用
llm.invoke(final_prompt.invoke(input = {"input":"1 || 10"}))


# 提示词的invoke输出给到了llm作为输入，和管道的概念一模一样
chain = final_prompt | llm
chain.invoke(input = {"input":"1 || 10"})