code: 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)
  
apis: langchain_visualizer.visualize langchain.llms.OpenAI langchain.chains.LLMChain langchain.PromptTemplate
extract_api: [((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')]