from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_models.tongyi import ChatTongyi

from dotenv import load_dotenv
from langchain_core.runnables import RunnableLambda

load_dotenv()

history_data = [
    ("human", "你来写一首诗？"),
    ("ai", "归园田居（其三） 种豆南山下，草盛豆苗稀。晨兴理荒秽，带月荷锄归。道狭草木长，夕露沾我衣。衣沾不足惜，但使愿无违。"),
    ("human", "好诗，再来一首"),
    ("ai", "饮酒（其五） 结庐在人境，而无车马喧。问君何能尔，心远地自偏。采菊东篱下，悠然见南山。山气日夕佳，飞鸟相与还。此中有真意，欲辨已忘言。")
]


# first_prompt_template = ChatPromptTemplate.from_messages([
#     ("system", "你是一个世外高人，不予世俗同流合污，以作诗表达心中所想,并封装成JSON格式返回给我。要求key是poetry，value是你作的诗"),
#     MessagesPlaceholder(variable_name="chat_history"),
#     ("human", "{input}"),
# ])

first_prompt_template = ChatPromptTemplate.from_messages([
    ("system", "你是一个世外高人，不予世俗同流合污，以作诗表达心中所感。"),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
])

second_prompt_template = ChatPromptTemplate.from_messages([
    ("ai", "上一首诗是：{poetry}"),
    ("human", "帮我解释一下诗名和每一句诗的意思"),
])

# 自定义函数
myfun = RunnableLambda(lambda ai_msg:{"poetry" : ai_msg.content}) #type: ignore

model = ChatTongyi(
    model = "qwen3-max"
)

str_parser = StrOutputParser()
json_parser = JsonOutputParser()
# 前一个输出等于下一个输入，必须是Runnable接口的子类，自定义函数入链需要遵守这个规则
chain = (first_prompt_template | model
         # json_parser
         | myfun
         | second_prompt_template | model | str_parser)

# 直接输出
# print(chain.invoke({"chat_history":history_data, "input": "再来一首诗"}).content)

# 流式输出
for chunk in chain.stream({"chat_history":history_data, "input": "再来一首诗"}):
    print(chunk, end="", flush=True)