import langchain_visualizer  # isort:skip  # noqa: F401
import asyncio

import vcr_langchain as vcr
from langchain import PromptTemplate
from langchain.chains import LLMChain
from langchain.llms import OpenAI

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


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

chain = LLMChain(llm=llm, prompt=prompt)


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


@vcr.use_cassette()
async def llm_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(llm_chain_demo())
    assert result.strip() == "Socktastic!"


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

    visualize(llm_chain_demo)

langchain.chains.LLMChain langchain_visualizer.visualize langchain.llms.OpenAI langchain.PromptTemplate
[((387, 408), 'langchain.llms.OpenAI', 'OpenAI', ([], {'temperature': '(0)'}), '(temperature=0)\n', (393, 408), False, 'from langchain.llms import OpenAI\n'), ((418, 534), '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", (432, 534), False, 'from langchain import PromptTemplate\n'), ((550, 582), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt'}), '(llm=llm, prompt=prompt)\n', (558, 582), False, 'from langchain.chains import LLMChain\n'), ((676, 694), 'vcr_langchain.use_cassette', 'vcr.use_cassette', ([], {}), '()\n', (692, 694), True, 'import vcr_langchain as vcr\n'), ((1042, 1067), 'langchain_visualizer.visualize', 'visualize', (['llm_chain_demo'], {}), '(llm_chain_demo)\n', (1051, 1067), False, 'from langchain_visualizer import visualize\n'), ((856, 880), 'asyncio.get_event_loop', 'asyncio.get_event_loop', ([], {}), '()\n', (878, 880), False, 'import asyncio\n')]