# Import OpenAI from langchain.llms
from langchain.llms import OpenAI

# Initialize OpenAI with model name and parameters
llm = OpenAI(model_name="text-ada-001", n=2, best_of=2)

# Generate a joke using the language model
llm("Tell me a joke")
# Output: "Why did the chicken cross the road? To get to the other side."

# Generate multiple outputs using the language model
llm_result = llm.generate(["Tell me a joke", "Tell me a poem"]*15)

# Print the length of the generations
print(len(llm_result.generations)) # Output: 30

# Print the first generation
print(llm_result.generations[0])

# Import PromptTemplate and OpenAI
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI

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

# Running the Chain
# Import LLMChain
from langchain.chains import LLMChain

# Initialize LLMChain with llm and prompt
chain = LLMChain(llm=llm, prompt=prompt)

# Execute the chain by specifying only the input variable.
print(chain.run("colorful socks"))

# Import ChatOpenAI, ChatPromptTemplate, and HumanMessagePromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)

# Initialize HumanMessagePromptTemplate with a prompt
human_message_prompt = HumanMessagePromptTemplate(
    prompt=PromptTemplate(
        template="What is a good name for a company that makes {product}

?",
        input_variables=["product"],
    )
)

# Initialize ChatPromptTemplate with messages
chat_prompt_template = ChatPromptTemplate.from_messages([human_message_prompt])

# Initialize ChatOpenAI
chat = ChatOpenAI(temperature=0.9)

# Initialize LLMChain with llm and prompt
chain = LLMChain(llm=chat, prompt=chat_prompt_template)

# Execute the chain by specifying only the input variable.
print(chain.run("colorful socks"))

# Initialize a second PromptTemplate
second_prompt = PromptTemplate(
    input_variables=["company_name"],
    template="Compose a catchphrase for the following company: {company_name}",
)

# Initialize a second LLMChain with llm and the second prompt
chain_two = LLMChain(llm=llm, prompt=second_prompt)

# Import SimpleSequentialChain
from langchain.chains import SimpleSequentialChain

# Initialize SimpleSequentialChain with chains and verbose mode
overall_chain = SimpleSequentialChain(chains=[chain, chain_two], verbose=True)

# Execute the chain by specifying only the input variable for the first chain.
catchphrase = overall_chain.run("colorful socks")

# Print the catchphrase
print(catchphrase)

# Import LLMChain, Chain, Dict, and List
from langchain.chains import LLMChain
from langchain.chains.base import Chain
from typing import Dict, List

# Define a new chain class: ConcatenateChain
class ConcatenateChain(Chain):
    def __init__(self, chain_1: LLMChain, chain_2: LLMChain):
        self.chain_1 = chain_1
        self.chain_2 = chain_2

    @property
    def input_keys(self) -> List[str]:
        return list(set(self.chain_1.input_keys).union(set(self.chain_2.input_keys)))

    def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
        output_1 = self.chain_1.run(inputs)
        output_2 = self.chain_2.run(inputs)
        return {'concat_output': output_1 + output_2}

# Initialize two PromptTemplates
prompt_1 = PromptTemplate(
    input_variables=["product"],
    template="What is a good name for a company that makes {product}?",
)
prompt_2 = PromptTemplate(
    input_variables=["product"],
    template="What is a catchy slogan for a company that makes {product}?",
)

# Initialize two LLMChains with llm and the prompts
chain_1 = LLMChain(llm=llm, prompt=prompt_1)
chain_2 = LLMChain(llm=llm, prompt=prompt_2)

# Initialize ConcatenateChain with the two chains
concat_chain = ConcatenateChain(chain_1=chain_1, chain_2=chain_2)

# Run the ConcatenateChain with an input
concat_output = concat_chain.run("colorful socks")

# Print the concatenated output
print(f"Concatenated output:\n{concat_output}")