code: import zipfile
  
  from langchain.chat_models import ChatOpenAI
  from langchain.schema import (
      HumanMessage,
      SystemMessage
  )
  
  import langchain
  from langchain.cache import SQLiteCache
  
  langchain.llm_cache = SQLiteCache(
      database_path=".langchain.db"
  )  # caches queries that are the same.
  
  def generate_code(question, model_type, api_key):
      """
      Generate Python code script. The script should only include code, no comments.
      What is returned is code that needs some light modification.
      """
      task = "Generate Python Code Script. The script should only include code, no comments."
      chat_model = ChatOpenAI(temperature=0, model=model_type, openai_api_key = api_key)
      response = chat_model(
          messages=[
              SystemMessage(content=task),
              HumanMessage(content=question)
          ],
      )
      code = response.content
      return format_code(code)
  
  
  def format_code(code):
      """
      Format the generated code to remove read_csv if its there.
      """
      csv_line_index = code.find("read_csv")
      if csv_line_index > 0:
          before_csv_line_index = code[0:csv_line_index].rfind("\n")
          before_csv = (
              "" if before_csv_line_index == -1 else code[0:before_csv_line_index]
          )
          after_csv = code[csv_line_index:]
          after_csv_line_index = after_csv.find("\n")
          after_csv = (
              "" if after_csv_line_index == -1 else after_csv[after_csv_line_index:]
          )
          return before_csv + after_csv
      return code
  
  
  def format_primer(primer_desc, primer_code, question):
      """
      Format the primer description and code.
      """
      return f'"""\n{primer_desc}{question}\n"""\n{primer_code}'
  
  
  def generate_primer(df, df_name):
      """
      Generate a primer for visualization. Include all the necessary imports.
      """
      primer_desc = """
      If a graph is needed, label the x and y axes appropriately.
      Add a title. Use st.plotly or plotly.express to visualize."
      Using Python version 3.9.12, only create a script using the dataframe "df" if the user asked for a graph. 
      Do not give a y/n answer in the code or try to take any user input. 
      Graph should be created using st.plotly or plotly.express. If there is only a single row, manipulate it to be able to chart the data.
      IF the given sql query is only one column, manipulate it to get it in the right format for a chart if it has multiple rows.
      If its just one column one row, don't chart, just display the table.
      Make sure that variable names are the same as the sql.
      """
      primer_code = (
          "import pandas as pd\nimport streamlit as st\nimport plotly.express as px\n"
      )
      primer_code += f"df = {df_name}.copy()\n"
      return primer_desc, primer_code
  
  def unzip(file):
      with zipfile.ZipFile(file, 'r') as zip_ref:
        zip_ref.extractall()
  
apis: langchain.schema.SystemMessage langchain.cache.SQLiteCache langchain.schema.HumanMessage langchain.chat_models.ChatOpenAI
extract_api: [((211, 253), 'langchain.cache.SQLiteCache', 'SQLiteCache', ([], {'database_path': '""".langchain.db"""'}), "(database_path='.langchain.db')\n", (222, 253), False, 'from langchain.cache import SQLiteCache\n'), ((621, 688), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0)', 'model': 'model_type', 'openai_api_key': 'api_key'}), '(temperature=0, model=model_type, openai_api_key=api_key)\n', (631, 688), False, 'from langchain.chat_models import ChatOpenAI\n'), ((2781, 2807), 'zipfile.ZipFile', 'zipfile.ZipFile', (['file', '"""r"""'], {}), "(file, 'r')\n", (2796, 2807), False, 'import zipfile\n'), ((749, 776), 'langchain.schema.SystemMessage', 'SystemMessage', ([], {'content': 'task'}), '(content=task)\n', (762, 776), False, 'from langchain.schema import HumanMessage, SystemMessage\n'), ((790, 820), 'langchain.schema.HumanMessage', 'HumanMessage', ([], {'content': 'question'}), '(content=question)\n', (802, 820), False, 'from langchain.schema import HumanMessage, SystemMessage\n')]