from langchain_openai import ChatOpenAI

model = ChatOpenAI(
    model_name="gpt-3.5-turbo",
    openai_api_key="sk-xxxx",
    openai_api_base="https://api.302.ai/v1",
)


from langchain.chains.router import LLMRouterChain, MultiPromptChain
from langchain.chains.router.llm_router import RouterOutputParser
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE
from langchain.chains import LLMChain, ConversationChain
from langchain.prompts import PromptTemplate

# 準(zhǔn)備2條目的鏈：一條物理鏈，一條數(shù)學(xué)鏈
# 1. 物理鏈
physics_template = """
你是一位物理學(xué)家，擅長(zhǎng)回答物理相關(guān)的問題，當(dāng)你不知道問題的答案時(shí)，你就回答不知道。
具體問題如下：
{input}
"""
physics_prompt = PromptTemplate.from_template(physics_template)
physics_chain = LLMChain(llm=model, prompt=physics_prompt)

# 2. 數(shù)學(xué)鏈
math_template = """
你是一個(gè)數(shù)學(xué)家，擅長(zhǎng)回答數(shù)學(xué)相關(guān)的問題，當(dāng)你不知道問題的答案時(shí)，你就回答不知道。
具體問題如下：
{input}
"""
math_prompt = PromptTemplate.from_template(math_template)
math_chain = LLMChain(llm=model, prompt=math_prompt)

# 3. 英語鏈
english_template = """
你是一個(gè)非常厲害的英語老師，擅長(zhǎng)回答英語相關(guān)的問題，當(dāng)你不知道問題的答案時(shí)，你就回答不知道。
具體問題如下：
{input}
"""
english_prompt = PromptTemplate.from_template(english_template)
english_chain = LLMChain(llm=model, prompt=english_prompt)


######### 所有可能的目的鏈
destination_chains = {}
destination_chains["physics"] = physics_chain
destination_chains["math"] = math_chain
destination_chains["english"] = english_chain


######### 默認(rèn)鏈
default_chain = ConversationChain(llm=model, output_key="text")

# 讓多路由模板 能找到合適的 提示詞模板
destinations_template_str = """
physics:擅長(zhǎng)回答物理問題
math:擅長(zhǎng)回答數(shù)學(xué)問題
english:擅長(zhǎng)回答英語問題
"""
router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(
    destinations=destinations_template_str
)

# 通過路由提示詞模板，構(gòu)建路由提示詞
router_prompt = PromptTemplate(
    template=router_template,
    input_variables=["input"],
    output_parser=RouterOutputParser(),
)

######### 路由鏈
router_chain = LLMRouterChain.from_llm(llm=model, prompt=router_prompt)

######### 最終的鏈
multi_prompt_chain = MultiPromptChain(
    router_chain=router_chain,
    destination_chains=destination_chains,
    default_chain=default_chain,
    verbose=True,
)



# multi_prompt_chain.invoke({"input": "重力加速度是多少？"})
# multi_prompt_chain.invoke("y=x^2+2x+1的導(dǎo)數(shù)是多少？")
multi_prompt_chain.invoke("將以下英文翻譯成中文，只輸出中文翻譯結(jié)果：\n The largest community building the future of LLM apps.")
# multi_prompt_chain.invoke("你是怎么理解java的面向?qū)ο蟮乃枷氲模?)