# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/gradio_application.ipynb.

# %% auto 0
__all__ = ['save_pdf', 'save_json', 'save_txt', 'save_csv', 'num_sources', 'css', 'save_chatbot_dialogue',
           'SlightlyDelusionalTutor', 'embed_key', 'create_reference_store', 'prompt_select', 'add_user_message',
           'get_tutor_reply', 'disable_until_done']

# %% ../nbs/gradio_application.ipynb 9
import gradio as gr
from functools import partial
import pandas as pd
import os

import PromptInteractionBase
import IOHelperUtilities
import SelfStudyPrompts
import MediaVectorStores

# %% ../nbs/gradio_application.ipynb 13
def save_chatbot_dialogue(chat_tutor, save_type):

    formatted_convo = pd.DataFrame(chat_tutor.conversation_memory, columns=['user', 'chatbot'])

    output_fname = f'tutoring_conversation.{save_type}'

    if save_type == 'csv':
        formatted_convo.to_csv(output_fname, index=False)
    elif save_type == 'json':
        formatted_convo.to_json(output_fname, orient='records')
    elif save_type == 'txt':
        temp = formatted_convo.apply(lambda x: 'User: {0}\nAI: {1}'.format(x[0], x[1]), axis=1)
        temp = '\n\n'.join(temp.tolist())
        with open(output_fname, 'w') as f:
            f.write(temp)
    else:
        gr.update(value=None, visible=False)
    
    return gr.update(value=output_fname, visible=True)

save_pdf = partial(save_chatbot_dialogue, save_type='pdf')
save_json = partial(save_chatbot_dialogue, save_type='json')
save_txt = partial(save_chatbot_dialogue, save_type='txt')
save_csv = partial(save_chatbot_dialogue, save_type='csv')


# %% ../nbs/gradio_application.ipynb 16
class SlightlyDelusionalTutor:
    # create basic initialization function
    def __init__(self, model_name = None):

        # create default model name
        if model_name is None:
            self.model_name = 'gpt-3.5-turbo-16k'

        self.chat_llm = None
        self.tutor_chain = None
        self.vector_store = None
        self.vs_retriever = None
        self.conversation_memory = []
        self.sources_memory = []
        self.flattened_conversation = ''
        self.api_key_valid = False
        self.learning_objectives = None
        self.openai_auth = ''
    
    def initialize_llm(self):

        if self.openai_auth:
            try:
                self.chat_llm = create_model(self.model_name, openai_api_key = self.openai_auth)
                self.api_key_valid = True
            except Exception as e:
                print(e)
                self.api_key_valid = False
        else:
            print("Please provide an OpenAI API key and press Enter.")

    def add_user_message(self, user_message):
        self.conversation_memory.append([user_message, None])
        self.flattened_conversation = self.flattened_conversation + '\n\n' + 'User: ' + user_message
    
    def get_tutor_reply(self, **input_kwargs):

        if not self.conversation_memory:
            return "Please type something to start the conversation."
        
        # we want to have a different vector comparison for reference lookup after the topic is first used
        if len(self.conversation_memory) > 1:
            if 'question' in input_kwargs.keys():
                if input_kwargs['question']:
                    input_kwargs['question'] = self.conversation_memory[-1][0] + ' keeping in mind I want to learn about ' + input_kwargs['question']
                else:
                    input_kwargs['question'] = self.conversation_memory[-1][0]

        # get tutor message
        tutor_message = get_tutoring_answer(None,
                                     self.tutor_chain,
                                     assessment_request = self.flattened_conversation + 'First, please provide your feedback on my previous answer if I was answering a question, otherwise, respond appropriately to my statement. Then, help me with the following:' + self.conversation_memory[-1][0],
                                     learning_objectives = self.learning_objectives,
                                     return_dict=True,
                                     **input_kwargs)

        # add tutor message to conversation memory
        self.conversation_memory[-1][1] = tutor_message['answer']
        self.flattened_conversation = self.flattened_conversation + '\nAI: ' + tutor_message['answer']
        self.sources_memory.append(tutor_message['source_documents'])
        #print(self.flattened_conversation, '\n\n')
        print(tutor_message['source_documents'])
    
    def get_sources_memory(self):
        # retrieve last source
        last_sources = self.sources_memory[-1]

        # get page_content keyword from last_sources
        doc_contents = ['Source ' + str(ind+1) + '\n"' + doc.page_content + '"\n\n' for ind, doc in enumerate(last_sources)]
        doc_contents = ''.join(doc_contents)

        return doc_contents
    
    def forget_conversation(self):
        self.conversation_memory = []
        self.sources_memory = []
        self.flattened_conversation = ''

# %% ../nbs/gradio_application.ipynb 18
def embed_key(openai_api_key, chat_tutor):
  if not openai_api_key:
    return chat_tutor
  
  # Otherwise, update key
  os.environ["OPENAI_API_KEY"] = openai_api_key
  
  #update tutor
  chat_tutor.openai_auth = openai_api_key

  if not chat_tutor.api_key_valid:
    chat_tutor.initialize_llm()

  return chat_tutor

# %% ../nbs/gradio_application.ipynb 20
def create_reference_store(chat_tutor, vs_button, text_cp, upload_files, reference_vs, openai_auth, learning_objs):

    text_segs = []
    upload_segs = []
    
    if reference_vs:
        raise NotImplementedError("Reference Vector Stores are not yet implemented")
    
    if text_cp.strip():
        text_segs = get_document_segments(text_cp, 'text', chunk_size=700, chunk_overlap=100)
        [doc.metadata.update({'source':'text box'}) for doc in text_segs];
    
    if upload_files:
        print(upload_files)
        upload_fnames = [f.name for f in upload_files]
        upload_segs = get_document_segments(upload_fnames, 'file', chunk_size=700, chunk_overlap=100)
    
    # get the full list of everything
    all_segs = text_segs + upload_segs
    print(all_segs)

    # create the vector store and update tutor
    vs_db, vs_retriever = create_local_vector_store(all_segs, search_kwargs={"k": 2})
    chat_tutor.vector_store = vs_db
    chat_tutor.vs_retriever = vs_retriever

    # create the tutor chain
    if not chat_tutor.api_key_valid or not chat_tutor.openai_auth:
        chat_tutor = embed_key(openai_auth, chat_tutor)
    qa_chain = create_tutor_mdl_chain(kind="retrieval_qa", mdl=chat_tutor.chat_llm, retriever = chat_tutor.vs_retriever, return_source_documents=True)
    chat_tutor.tutor_chain = qa_chain

    # store learning objectives
    chat_tutor.learning_objectives = learning_objs

    # return the story
    return chat_tutor, gr.update(interactive=True, value='Tutor Initialized!')

# %% ../nbs/gradio_application.ipynb 22
### Gradio Called Functions ###

def prompt_select(selection, number, length):
  if selection == "Random":
    prompt = f"Please design a {number} question quiz based on the context provided and the inputted learning objectives (if applicable). The types of questions should be randomized (including multiple choice, short answer, true/false, short answer, etc.).  Provide one question at a time, and wait for my response before providing me with feedback. Again, while the quiz may ask for multiple questions, you should only provide 1 question in you initial response. Do not include the answer in your response. If I get an answer wrong, provide me with an explanation of why it was incorrect, and then give me additional chances to respond until I get the correct choice. Explain why the correct choice is right."
  elif selection == "Fill in the Blank":
    prompt = f"Create a {number} question fill in the blank quiz refrencing the context provided. The quiz should reflect the learning objectives (if inputted). The 'blank' part of the question should appear as '________'. The answers should reflect what word(s) should go in the blank an accurate statement. An example is the follow: 'The author of the article is ______.' The question should be a statement. Provide one question at a time, and wait for my response before providing me with feedback. Again, while the quiz may ask for multiple questions, you should only provide ONE question in you initial response. Do not include the answer in your response. If I get an answer wrong, provide me with an explanation of why it was incorrect,and then give me additional chances to respond until I get the correct choice. Explain why the correct choice is right."
  elif selection == "Short Answer":
    prompt = f"Please design a {number} question quiz about which reflects the learning objectives (if inputted).  The questions should be short answer. Expect the correct answers to be {length} sentences long. Provide one question at a time, and wait for my response before providing me with feedback. Again, while the quiz may ask for multiple questions, you should only provide ONE question in you initial response. Do not include the answer in your response. If I get an answer wrong, provide me with an explanation of why it was incorrect, and then give me additional chances to respond until I get the correct choice. Explain why the correct answer is right."
  else:
    prompt = f"Please design a {number} question {selection.lower()} quiz based on the context provided and the inputted learning objectives (if applicable). Provide one question at a time, and wait for my response before providing me with feedback. Again, while the quiz may ask for multiple questions, you should only provide 1 question in you initial response. Do not include the answer in your response. If I get an answer wrong, provide me with an explanation of why it was incorrect, and then give me additional chances to respond until I get the correct choice. Explain why the correct choice is right."
  return prompt, prompt
    

# %% ../nbs/gradio_application.ipynb 24
### Chatbot Functions ###

def add_user_message(user_message, chat_tutor):
  """Display user message and update chat history to include it.
  Also disables user text input until bot is finished (call to reenable_chat())
  See https://gradio.app/creating-a-chatbot/"""
  chat_tutor.add_user_message(user_message)
  return gr.update(value="", interactive=False), chat_tutor.conversation_memory, chat_tutor

def get_tutor_reply(learning_topic, chat_tutor):
  chat_tutor.get_tutor_reply(input_kwargs={'question':learning_topic})
  return gr.update(value="", interactive=True), gr.update(visible=True, value=chat_tutor.get_sources_memory()), chat_tutor.conversation_memory, chat_tutor

num_sources = 2

# %% ../nbs/gradio_application.ipynb 25
def disable_until_done(obj_in):
    return gr.update(interactive=False)

# %% ../nbs/gradio_application.ipynb 27
# See https://gradio.app/custom-CSS-and-JS/
css="""
#sources-container {
  overflow: scroll !important; /* Needs to override default formatting */
  /*max-height: 20em; */ /* Arbitrary value */
}
#sources-container > div { padding-bottom: 1em !important; /* Arbitrary value */ }
.short-height > * > * { min-height: 0 !important; }
.translucent { opacity: 0.5; }
.textbox_label { padding-bottom: .5em; }
"""
#srcs = [] # Reset sources (db and qa are kept the same for ease of testing)

with gr.Blocks(css=css, analytics_enabled=False) as demo:

  #initialize tutor (with state)
  study_tutor = gr.State(SlightlyDelusionalTutor())

  # Title
  gr.Markdown("# Studying with a Slightly Delusional Tutor")

  # API Authentication functionality
  with gr.Box():
    gr.Markdown("### OpenAI API Key ")
    gr.HTML("""<span>Embed your OpenAI API key below; if you haven't created one already, visit
      <a href="https://platform.openai.com/account/api-keys">platform.openai.com/account/api-keys</a>
    to sign up for an account and get your personal API key</span>""",
            elem_classes="textbox_label")
    api_input = gr.Textbox(show_label=False, type="password", container=False, autofocus=True,
                      placeholder="●●●●●●●●●●●●●●●●●", value='')
    api_input.submit(fn=embed_key, inputs=[api_input, study_tutor], outputs=study_tutor)
    api_input.blur(fn=embed_key, inputs=[api_input, study_tutor], outputs=study_tutor)

  # Reference document functionality (building vector stores)
  with gr.Box():
    gr.Markdown("### Add Reference Documents")
    # TODO Add entry for path to vector store (should be disabled for now)
    with gr.Row(equal_height=True):
      text_input = gr.TextArea(label='Copy and paste your text below',
                                lines=2)

      file_input = gr.Files(label="Load a .txt or .pdf file",
                            file_types=['.pdf', '.txt'], type="file",
                            elem_classes="short-height")

      instructor_input = gr.TextArea(label='Enter vector store URL, if given by instructor (WIP)', value='',
                                lines=2, interactive=False, elem_classes="translucent")
    
    # Adding the learning objectives
    with gr.Box():
      gr.Markdown("### Optional: Enter Your Learning Objectives")
      learning_objectives = gr.Textbox(label='If provided by your instructor, please input your learning objectives for this session', value='')
    
    # Adding the button to submit all of the settings and create the Chat Tutor Chain.
    with gr.Row():
      vs_build_button = gr.Button(value = 'Start Studying with Your Tutor!', scale=1)
      vs_build_button.click(disable_until_done, vs_build_button, vs_build_button) \
        .then(create_reference_store, [study_tutor, vs_build_button, text_input, file_input, instructor_input, api_input, learning_objectives],
                            [study_tutor, vs_build_button])

  

  # Premade question prompts
  with gr.Box():
    gr.Markdown("""
    ## Generate a Premade Prompt
    Select your type and number of desired questions. Click "Generate Prompt" to get your premade prompt,
    and then "Insert Prompt into Chat" to copy the text into the chat interface below. \
    You can also copy the prompt using the icon in the upper right corner and paste directly into the input box when interacting with the model.
    """)
    with gr.Row():
      with gr.Column():
        question_type = gr.Dropdown(["Multiple Choice", "True or False", "Short Answer", "Fill in the Blank", "Random"], label="Question Type")
        number_of_questions = gr.Textbox(label="Enter desired number of questions")
        sa_desired_length = gr.Dropdown(["1-2", "3-4", "5-6", "6 or more"], label = "For short answer questions only, choose the desired sentence length for answers. The default value is 1-2 sentences.")
      with gr.Column():
        prompt_button = gr.Button("Generate Prompt")
        premade_prompt_output = gr.Textbox(label="Generated prompt (save or copy)", show_copy_button=True)


  # Chatbot interface
  gr.Markdown("## Chat with the Model")
  topic_input = gr.Textbox(label="What topic or concept are you trying to learn more about?")
  with gr.Row(equal_height=True):
    with gr.Column(scale=2):
      chatbot = gr.Chatbot()
      with gr.Row():
        user_chat_input = gr.Textbox(label="User input", scale=9)
        user_chat_submit = gr.Button("Ask/answer model", scale=1)

    # sources
    with gr.Box(elem_id="sources-container", scale=1):
      # TODO: Display document sources in a nicer format?
      gr.HTML(value="<h3 id='sources'>Referenced Sources</h3>")
      sources_output = gr.Textbox(value='', interactive=False, visible=False, show_label=False)
      #sources_output = []
      #for i in range(num_sources):
      #  source_elem = gr.HTML(visible=False)
      #  sources_output.append(source_elem)

  #define the behavior of prompt button later since it depends on user_chat_input
  prompt_button.click(prompt_select,
                      inputs=[question_type, number_of_questions, sa_desired_length],
                      outputs=[premade_prompt_output, user_chat_input])

  # Display input and output in three-ish parts
  # (using asynchronous functions):
  # First show user input, then show model output when complete
  # Then wait until the bot provides response and return the result
  # Finally, allow the user to ask a new question by reenabling input
  async_response = user_chat_submit.click(add_user_message,
            [user_chat_input, study_tutor],
            [user_chat_input, chatbot, study_tutor], queue=False) \
  .then(get_tutor_reply, [topic_input, study_tutor], [user_chat_input, sources_output, chatbot, study_tutor], queue=True)

  async_response_b = user_chat_input.submit(add_user_message,
            [user_chat_input, study_tutor],
            [user_chat_input, chatbot, study_tutor], queue=False) \
  .then(get_tutor_reply, [topic_input, study_tutor], [user_chat_input, sources_output, chatbot, study_tutor], queue=True)

  with gr.Blocks():
    gr.Markdown("""
    ## Export Your Chat History
    Export your chat history as a .json, PDF file, .txt, or .csv file
    """)
    with gr.Row():
      export_dialogue_button_json = gr.Button("JSON")
      export_dialogue_button_pdf = gr.Button("PDF")
      export_dialogue_button_txt = gr.Button("TXT")
      export_dialogue_button_csv = gr.Button("CSV")
    
    file_download = gr.Files(label="Download here",
                            file_types=['.pdf', '.txt', '.csv', '.json'], type="file", visible=False)
    
    export_dialogue_button_json.click(save_json, study_tutor, file_download, show_progress=True)
    export_dialogue_button_pdf.click(save_pdf, study_tutor, file_download, show_progress=True)
    export_dialogue_button_txt.click(save_txt, study_tutor, file_download, show_progress=True)
    export_dialogue_button_csv.click(save_csv, study_tutor, file_download, show_progress=True)

demo.queue()
demo.launch(debug=True)
#demo.launch()
#gr.close_all()