code: import os
  from typing import Optional
  
  from langchain import LLMChain, OpenAI, PromptTemplate
  from langchain.chains.base import Chain
  from langchain.chains.conversation.memory import ConversationBufferMemory
  from langchain.llms.base import BaseLLM
  from langchain.llms.loading import load_llm
  
  
  DEFAULT_LLM = None
  # Default template, no memory
  TEMPLATE = """
  You are working with a pandas dataframe in Python. The name of the dataframe is `df`.
  The dataframe has the following columns: {df_columns}.
  
  You should execute code as commanded to either provide information to answer the question or to
  do the transformations required.
  
  You should not assign any variables; you should return a one-liner in Pandas.
  
  This is your objective: {query}
  
  Go!
  
  ```python
  print(df.head())
  ```
  ```output
  {df_head}
  ```
  ```python"""
  
  PROMPT = PromptTemplate(template=TEMPLATE, input_variables=["query", "df_head", "df_columns"])
  
  
  # Template with memory
  # TODO: add result of expected code to memory; currently we only remember what code was run.
  TEMPLATE_WITH_MEMORY = """
  You are working with a pandas dataframe in Python. The name of the dataframe is `df`.
  The dataframe has the following columns: {df_columns}.
  
  You are interacting with a programmer. The programmer issues commands and you should translate
  them into Python code and execute them.
  
  This is the history of your interaction so far:
  {chat_history}
  Human: {query}
  
  Go!
  
  ```python
  df.head()
  ```
  ```output
  {df_head}
  ```
  ```python
  """
  PROMPT_WITH_MEMORY = PromptTemplate(
      template=TEMPLATE_WITH_MEMORY, input_variables=["chat_history", "query", "df_head", "df_columns"]
  )
  
  
  def set_llm(llm: BaseLLM) -> None:
      global DEFAULT_LLM
      DEFAULT_LLM = llm
  
  
  def get_chain(llm: Optional[BaseLLM] = None, use_memory: bool = True) -> Chain:
      """Get chain to use."""
      if llm is None:
          if DEFAULT_LLM is None:
              llm_config_path = os.environ.get("LLPANDAS_LLM_CONFIGURATION")
              if llm_config_path is None:
                  llm = OpenAI(temperature=0)
              else:
                  llm = load_llm(llm_config_path)
          else:
              llm = DEFAULT_LLM
  
      if use_memory:
          memory = ConversationBufferMemory(memory_key="chat_history", input_key="query")
          chain = LLMChain(llm=llm, prompt=PROMPT_WITH_MEMORY, memory=memory)
      else:
          chain = LLMChain(llm=llm, prompt=PROMPT)
  
      return chain
  
apis: langchain.PromptTemplate langchain.chains.conversation.memory.ConversationBufferMemory langchain.llms.loading.load_llm langchain.LLMChain langchain.OpenAI
extract_api: [((825, 914), 'langchain.PromptTemplate', 'PromptTemplate', ([], {'template': 'TEMPLATE', 'input_variables': "['query', 'df_head', 'df_columns']"}), "(template=TEMPLATE, input_variables=['query', 'df_head',\n    'df_columns'])\n", (839, 914), False, 'from langchain import LLMChain, OpenAI, PromptTemplate\n'), ((1501, 1619), 'langchain.PromptTemplate', 'PromptTemplate', ([], {'template': 'TEMPLATE_WITH_MEMORY', 'input_variables': "['chat_history', 'query', 'df_head', 'df_columns']"}), "(template=TEMPLATE_WITH_MEMORY, input_variables=[\n    'chat_history', 'query', 'df_head', 'df_columns'])\n", (1515, 1619), False, 'from langchain import LLMChain, OpenAI, PromptTemplate\n'), ((2171, 2241), 'langchain.chains.conversation.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'input_key': '"""query"""'}), "(memory_key='chat_history', input_key='query')\n", (2195, 2241), False, 'from langchain.chains.conversation.memory import ConversationBufferMemory\n'), ((2258, 2317), 'langchain.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'PROMPT_WITH_MEMORY', 'memory': 'memory'}), '(llm=llm, prompt=PROMPT_WITH_MEMORY, memory=memory)\n', (2266, 2317), False, 'from langchain import LLMChain, OpenAI, PromptTemplate\n'), ((2344, 2376), 'langchain.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'PROMPT'}), '(llm=llm, prompt=PROMPT)\n', (2352, 2376), False, 'from langchain import LLMChain, OpenAI, PromptTemplate\n'), ((1895, 1939), 'os.environ.get', 'os.environ.get', (['"""LLPANDAS_LLM_CONFIGURATION"""'], {}), "('LLPANDAS_LLM_CONFIGURATION')\n", (1909, 1939), False, 'import os\n'), ((2002, 2023), 'langchain.OpenAI', 'OpenAI', ([], {'temperature': '(0)'}), '(temperature=0)\n', (2008, 2023), False, 'from langchain import LLMChain, OpenAI, PromptTemplate\n'), ((2064, 2089), 'langchain.llms.loading.load_llm', 'load_llm', (['llm_config_path'], {}), '(llm_config_path)\n', (2072, 2089), False, 'from langchain.llms.loading import load_llm\n')]