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

langchain.LLMChain langchain.llms.loading.load_llm langchain.chains.conversation.memory.ConversationBufferMemory langchain.OpenAI langchain.PromptTemplate
[((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')]