from langchain.chains.router import MultiPromptChain
from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser
from langchain.prompts import PromptTemplate

from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.chains import LLMChain
from app.type import ChatGPTModel
import logging

physics_template = """You are a very smart physics professor. \
        You are great at answering questions about physics in a concise\
        and easy to understand manner. \
        When you don't know the answer to a question you admit\
        that you don't know.

        Here is a question:
        {input}"""

math_template = """You are a very good mathematician. \
        You are great at answering math questions. \
        You are so good because you are able to break down \
        hard problems into their component parts, 
        answer the component parts, and then put them together\
        to answer the broader question.

        Here is a question:
        {input}"""

history_template = """You are a very good historian. \
        You have an excellent knowledge of and understanding of people,\
        events and contexts from a range of historical periods. \
        You have the ability to think, reflect, debate, discuss and \
        evaluate the past. You have a respect for historical evidence\
        and the ability to make use of it to support your explanations \
        and judgements.

        Here is a question:
        {input}"""

computerscience_template = """ You are a successful computer scientist.\
        You have a passion for creativity, collaboration,\
        forward-thinking, confidence, strong problem-solving capabilities,\
        understanding of theories and algorithms, and excellent communication \
        skills. You are great at answering coding questions. \
        You are so good because you know how to solve a problem by \
        describing the solution in imperative steps \
        that a machine can easily interpret and you know how to \
        choose a solution that has a good balance between \
        time complexity and space complexity. 

        Here is a question:
        {input}"""

prompt_infos = [
    {
        "name": "physics",
        "description": "Good for answering questions about physics",
        "prompt_template": physics_template
    },
    {
        "name": "math",
        "description": "Good for answering math questions",
        "prompt_template": math_template
    },
    {
        "name": "History",
        "description": "Good for answering history questions",
        "prompt_template": history_template
    },
    {
        "name": "computer science",
        "description": "Good for answering computer science questions",
        "prompt_template": computerscience_template
    }
]

MULTI_PROMPT_ROUTER_TEMPLATE = """Given a raw text input to a \
language model select the model prompt best suited for the input. \
You will be given the names of the available prompts and a \
description of what the prompt is best suited for. \
You may also revise the original input if you think that revising\
it will ultimately lead to a better response from the language model.

<< FORMATTING >>
Return a markdown code snippet with a JSON object formatted to look like:
```json
{{{{
    "destination": string \ name of the prompt to use or "DEFAULT"
    "next_inputs": string \ a potentially modified version of the original input
}}}}
```

REMEMBER: "destination" MUST be one of the candidate prompt \
names specified below OR it can be "DEFAULT" if the input is not\
well suited for any of the candidate prompts.
REMEMBER: "next_inputs" can just be the original input \
if you don't think any modifications are needed.

<< CANDIDATE PROMPTS >>
{destinations}

<< INPUT >>
{{input}}

<< OUTPUT (remember to include the ```json)>>"""


class MultiChain:
    chain: any

    def __init__(self):
        llm = ChatOpenAI(temperature=0, model=ChatGPTModel.GPT3.value)

        destination_chains = {}
        for p_info in prompt_infos:
            name = p_info["name"]
            prompt_template = p_info["prompt_template"]
            prompt = ChatPromptTemplate.from_template(template=prompt_template)
            chain = LLMChain(llm=llm, prompt=prompt)
            destination_chains[name] = chain

        destinations = [f"{p['name']}: {p['description']}" for p in prompt_infos]
        destinations_str = "\n".join(destinations)

        default_prompt = ChatPromptTemplate.from_template("{input}")
        default_chain = LLMChain(llm=llm, prompt=default_prompt)

        router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(
            destinations=destinations_str
        )
        router_prompt = PromptTemplate(
            template=router_template,
            input_variables=["input"],
            output_parser=RouterOutputParser(),
        )

        router_chain = LLMRouterChain.from_llm(llm, router_prompt)

        self.chain = MultiPromptChain(router_chain=router_chain,
                                      destination_chains=destination_chains,
                                      default_chain=default_chain, verbose=True
                                      )

    def chain_query(self, input):
        logging.info(f"执行Chain查询,输入{input}")
        result = self.chain.run(input=input)
        logging.info(f"执行Chain查询,输入{input}，输出{result}")
        return result


if __name__ == "__main__":
    import langchain

    langchain.debug = True
    chain = MultiChain()
    result = chain.chain_query("中国最早的朝代")
    print(result)

langchain.chains.router.llm_router.LLMRouterChain.from_llm langchain.chat_models.ChatOpenAI langchain.chains.router.MultiPromptChain langchain.chains.router.llm_router.RouterOutputParser langchain.chains.LLMChain langchain.prompts.ChatPromptTemplate.from_template
[((3977, 4033), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0)', 'model': 'ChatGPTModel.GPT3.value'}), '(temperature=0, model=ChatGPTModel.GPT3.value)\n', (3987, 4033), False, 'from langchain.chat_models import ChatOpenAI\n'), ((4531, 4574), 'langchain.prompts.ChatPromptTemplate.from_template', 'ChatPromptTemplate.from_template', (['"""{input}"""'], {}), "('{input}')\n", (4563, 4574), False, 'from langchain.prompts import ChatPromptTemplate\n'), ((4599, 4639), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'default_prompt'}), '(llm=llm, prompt=default_prompt)\n', (4607, 4639), False, 'from langchain.chains import LLMChain\n'), ((4955, 4998), 'langchain.chains.router.llm_router.LLMRouterChain.from_llm', 'LLMRouterChain.from_llm', (['llm', 'router_prompt'], {}), '(llm, router_prompt)\n', (4978, 4998), False, 'from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser\n'), ((5021, 5151), 'langchain.chains.router.MultiPromptChain', 'MultiPromptChain', ([], {'router_chain': 'router_chain', 'destination_chains': 'destination_chains', 'default_chain': 'default_chain', 'verbose': '(True)'}), '(router_chain=router_chain, destination_chains=\n    destination_chains, default_chain=default_chain, verbose=True)\n', (5037, 5151), False, 'from langchain.chains.router import MultiPromptChain\n'), ((5305, 5341), 'logging.info', 'logging.info', (['f"""执行Chain查询,输入{input}"""'], {}), "(f'执行Chain查询,输入{input}')\n", (5317, 5341), False, 'import logging\n'), ((5395, 5442), 'logging.info', 'logging.info', (['f"""执行Chain查询,输入{input}，输出{result}"""'], {}), "(f'执行Chain查询,输入{input}，输出{result}')\n", (5407, 5442), False, 'import logging\n'), ((4214, 4272), 'langchain.prompts.ChatPromptTemplate.from_template', 'ChatPromptTemplate.from_template', ([], {'template': 'prompt_template'}), '(template=prompt_template)\n', (4246, 4272), False, 'from langchain.prompts import ChatPromptTemplate\n'), ((4293, 4325), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt'}), '(llm=llm, prompt=prompt)\n', (4301, 4325), False, 'from langchain.chains import LLMChain\n'), ((4899, 4919), 'langchain.chains.router.llm_router.RouterOutputParser', 'RouterOutputParser', ([], {}), '()\n', (4917, 4919), False, 'from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser\n')]