| from langchain import PromptTemplate, LLMChain | |
| from langchain.llms import OpenAI | |
| import openai | |
| import os | |
| template = """You are given following information and a question. Generate a python code with plotly to find the answer. Provide Text as comments only in code. | |
| Information: | |
| {main_question} | |
| ============= | |
| Question: | |
| {sub_question} | |
| Answer:""" | |
| prompt = PromptTemplate(template=template, input_variables=["main_question", "sub_question"]) | |
| def set_openai_api_key(api_key): | |
| openai.api_key = api_key | |
| def generate_plot(main_question, sub_question): | |
| llm = OpenAI(temperature=0) | |
| llm_chain = LLMChain(prompt=prompt, llm=llm) | |
| response = llm_chain.run(main_question=main_question, sub_question=sub_question) | |
| return response |