import langchain_visualizer  # isort:skip  # noqa: F401
import asyncio
from typing import Any, Dict, List, Optional

import vcr_langchain as vcr
from langchain import PromptTemplate
from langchain.callbacks.manager import CallbackManagerForChainRun
from langchain.chains import LLMChain
from langchain.chains.base import Chain
from langchain.llms import OpenAI

# ========================== Start of langchain example code ==========================
# https://langchain.readthedocs.io/en/latest/modules/chains/getting_started.html


class ConcatenateChain(Chain):
    chain_1: LLMChain
    chain_2: LLMChain

    @property
    def input_keys(self) -> List[str]:
        # Union of the input keys of the two chains.
        all_input_vars = set(self.chain_1.input_keys).union(
            set(self.chain_2.input_keys)
        )
        return list(all_input_vars)

    @property
    def output_keys(self) -> List[str]:
        return ["concat_output"]

    def _call(
        self,
        inputs: Dict[str, str],
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:
        output_1 = self.chain_1.run(inputs)
        output_2 = self.chain_2.run(inputs)
        return {"concat_output": output_1 + output_2}


llm = OpenAI()

prompt_1 = PromptTemplate(
    input_variables=["product"],
    template="What is a good name for a company that makes {product}?",
)
chain_1 = LLMChain(llm=llm, prompt=prompt_1)

prompt_2 = PromptTemplate(
    input_variables=["product"],
    template="What is a good slogan for a company that makes {product}?",
)
chain_2 = LLMChain(llm=llm, prompt=prompt_2)

concat_chain = ConcatenateChain(chain_1=chain_1, chain_2=chain_2)
chain = concat_chain


# ================================== Execute example ==================================


@vcr.use_cassette()
async def custom_chain_demo():
    return chain.run("colorful socks")


def test_llm_usage_succeeds():
    """Check that the chain can run normally"""
    result = asyncio.get_event_loop().run_until_complete(custom_chain_demo())
    assert (
        result.strip()
        == 'Sock Spectacular.\n\n"Step Up Your Style with Colorful Socks!"'
    )


if __name__ == "__main__":
    from langchain_visualizer import visualize

    visualize(custom_chain_demo)

langchain.chains.LLMChain langchain_visualizer.visualize langchain.llms.OpenAI langchain.PromptTemplate
[((1254, 1262), 'langchain.llms.OpenAI', 'OpenAI', ([], {}), '()\n', (1260, 1262), False, 'from langchain.llms import OpenAI\n'), ((1275, 1391), 'langchain.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['product']", 'template': '"""What is a good name for a company that makes {product}?"""'}), "(input_variables=['product'], template=\n    'What is a good name for a company that makes {product}?')\n", (1289, 1391), False, 'from langchain import PromptTemplate\n'), ((1408, 1442), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt_1'}), '(llm=llm, prompt=prompt_1)\n', (1416, 1442), False, 'from langchain.chains import LLMChain\n'), ((1455, 1573), 'langchain.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['product']", 'template': '"""What is a good slogan for a company that makes {product}?"""'}), "(input_variables=['product'], template=\n    'What is a good slogan for a company that makes {product}?')\n", (1469, 1573), False, 'from langchain import PromptTemplate\n'), ((1590, 1624), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt_2'}), '(llm=llm, prompt=prompt_2)\n', (1598, 1624), False, 'from langchain.chains import LLMChain\n'), ((1806, 1824), 'vcr_langchain.use_cassette', 'vcr.use_cassette', ([], {}), '()\n', (1822, 1824), True, 'import vcr_langchain as vcr\n'), ((2253, 2281), 'langchain_visualizer.visualize', 'visualize', (['custom_chain_demo'], {}), '(custom_chain_demo)\n', (2262, 2281), False, 'from langchain_visualizer import visualize\n'), ((1989, 2013), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (2011, 2013), False, 'import asyncio\n')]