functions.py

import os
import csv
import time
import openai
import gradio as gr


def get_countries():
    """
    Returns a list with the countries from a CSV
    """
    with open('assets/data/paises.csv', encoding='utf-8') as file:
        reader = csv.reader(file)
        countries = [row[0] for row in reader]

    return countries


def check_info(name: gr.Textbox, age: gr.Number, country: gr.Dropdown):
    """
    Checks if the given info is complete (no empty strings) and if the
    user is old enough (> 18). If so, opens the chatbot, if not
    """

    # Name or country is not provided or the age is not an integer
    if name == '' or country == '' or not isinstance(age, float):
        new_image = 'assets/images/pregunta.jpg'
        keep_asking, keep_open, open_chat = True, True, False
        msg_history = []

    # User is underage
    elif age < 18:
        new_image = 'assets/images/despedida.jpg'
        keep_asking, keep_open, open_chat = False, True, False
        msg_history = []

    # User is old enough and gave all the necessary data
    else:
        new_image = 'assets/images/feliz.jpg'
        keep_asking, keep_open, open_chat = False, False, True
        msg_history = innit_bot()

    # Update layout depending on the previous values
    new_button = gr.Button.update(interactive=keep_asking)
    new_img = gr.Image.update(value=new_image)
    new_welcoming = gr.Row.update(visible=keep_open)
    new_chat = gr.Row.update(visible=open_chat)

    return new_button, new_img, new_welcoming, new_chat, msg_history


def innit_bot():
    """
    Initialize the bot by adding the prompt from the txt file to the messages history
    """
    openai.api_key = os.environ.get('API_KEY')

    with open('assets/data/prompt.txt', encoding='utf-8') as file:
        prompt = file.read()
    message_history = [{"role": "system", "content": prompt}]

    return message_history


def call_api(msg_history: gr.State):
    """
    Returns the APIs response
    """
    response = openai.ChatCompletion.create(
        model="gpt-4",
        temperature=1,
        messages=msg_history
    )
    return response


def handle_call(msg_history: gr.State):
    """
    Returns the status of the response (False if there was an error, True otherwise), the response and
    waiting time of the AI. It also handles the possible errors
    """
    tries, max_tries = 0, 9
    while True:
        try:
            start_time = time.time()
            response = call_api(msg_history)
            end_time = time.time()
            break

        except Exception as e:
            print(e)

            if tries == max_tries:
                return True, '', ''

            tries += 1
            time.sleep(20)

    needed_time = end_time - start_time
    return False, response, needed_time


def get_ai_answer(msg: str, msg_history: gr.State):
    """
    Returns the response given by the model, all the message history so far and the seconds
    the api took to retrieve such response. It can also return an empty string if there was an
    error
    """
    msg_history.append({"role": "user", "content": msg})
    error, response, needed_time = handle_call(msg_history)

    # There was an error with the API, abort everything
    if error:
        return '', msg_history, -1

    AI_response = response["choices"][0]["message"]["content"]
    msg_history.append({'role': 'assistant', 'content': AI_response})

    return AI_response, msg_history, needed_time


def get_answer(msg: str, msg_history: gr.State, chatbot_history: gr.Chatbot, waiting_time: gr.State):
    """
    Cleans msg box, adds the new message to the message history,
    gets the answer from the bot and adds it to the chatbot history
    and gets the time needed to get such answer and saves it
    """

    # Get bot answer (output), messages history and waiting time
    AI_response, msg_history, needed_time = get_ai_answer(msg, msg_history)

    # There was an error
    if AI_response == '':
        waiting_time.append(needed_time)
        chatbot_history.append((msg, ''))
        return "", msg_history, chatbot_history, waiting_time, False

    # Make sure the AI_response is short, if not make it shorter
    if len(AI_response) > 260:
        new_msg = 'El mensaje esta muy largo. Da la misma idea (mandando el link, pregunta y/o promocion que hayas ' \
                  'dado) pero usando 40 palabras.'
        AI_response, msg_history, needed_time = get_ai_answer(new_msg, msg_history)

    # Make sure the AI_response has at least one question
    if '?' not in AI_response:
        new_msg = 'Incluye 1 pregunta dentro del mensaje. Puede estar relacionada a lo que se hablo antes o algo nuevo.'
        AI_response, msg_history, needed_time = get_ai_answer(new_msg, msg_history)

    # Save waiting time
    waiting_time.append(needed_time)

    # Save output in the chat
    chatbot_history.append((msg, AI_response))

    return "", msg_history, chatbot_history, waiting_time, True


def check_closing(is_working: gr.State):
    """
    If is_working is False (it means the API had an error), then close the chat
    """
    if is_working:
        return gr.Column.update(visible=True), gr.Column.update(visible=False)
    return gr.Column.update(visible=False), gr.Column.update(visible=True)