import os
import openai
from llama_index import StorageContext, load_index_from_storage, LLMPredictor, ServiceContext
from langchain.chat_models import ChatOpenAI
import gradio as gr

# Conectar Cuenta API de OpenAI
openai_api_key = os.environ.get('openai_key')
if openai_api_key:
    os.environ["OPENAI_API_KEY"] = openai_api_key
    openai.api_key = openai_api_key
else:
    print("Error con la clave de acceso a OpenAI.")

# Cargar entrenamiento y Modelo
exec(os.environ.get('storage_context'))

prompt = 'responder en español como un asistente experto en medicina, dando una respuesta detallada y reflejando de forma fiel los datos disponibles. Considerar segun aplique la siguiente informacion adicional: Nombre del Medicamento: Acsodix. Principio Activo: Vortioxetina. Formas y Dosis: Acsodix 5, 10, y 20 mg (comprimidos recubiertos). Indicaciones: Tratamiento de episodios de depresión mayor en adultos. Contraindicaciones: Alergia a la Vortioxetina, toma de ciertos inhibidores de la monoaminooxidasa, enfermedad hepática grave, etc. Precauciones: Embarazo, lactancia, conducción, uso de otros medicamentos, etc. Efectos Adversos: Náuseas, diarrea, vómitos, mareo, prurito, etc. Instrucciones de Uso: Se puede tomar con o sin alimentos, no se aconseja la combinación con alcohol, etc. Presentación: Cajas conteniendo 30 comprimidos recubiertos en diferentes dosis. La dosis máxima de Acsodix para adultos menores a 65 años es de 20 mg al día. Se debe tener precaución cuando se traten pacientes mayores de 65 años de edad con dosis superiores a 10 mg de Vortioxetina una vez al día, ya que los datos son limitados. El medicamento Vortioxetina (Acsodix) no puede ni debe ser combinado con alcohol. Se debe tener precaución al combinar la Vortioxetina con anticoagulantes o antiagregantes plaquetarios orales como warfarina, dipiridamol, fenprocumón, ácido acetilsalicílico, debido al potencial aumento del riesgo de hemorragia. La información proporcionada no menciona específicamente los dolores de cabeza como un efecto secundario de tomar Acsodix. Informe a su médico si está tomando, ha tomado recientemente o podría tener que tomar cualquier otro medicamento, incluso los adquiridos sin receta. No se observaron cambios en los niveles de hormonas sexuales después de la administración conjunta de Vortioxetina con el anticonceptivo oral combinado (etinil estradiol 30 μg/ levonorgestrel 150 μg).Acsodix no está recomendado en niños y adolescentes menores de 18 años, debido a la falta de información en este grupo de edad. La pregunta a responder es la siguiente:'

# Funcion de chat
def predict(message):
    # Agrega la pregunta del usuario al historial en formato OpenAI
    history_openai_format = [{"role": "user", "content": message}]
    
    # Agrega el prompt personalizado al inicio del historial
    full_prompt = [{'role': 'assistant', 'content': prompt}]
    full_prompt.extend(history_openai_format)
    
    # Obtener la pregunta del usuario
    user_message = full_prompt[-1]["content"]
    
    # Utilizar el índice de entrenamiento para generar la respuesta
    query = user_message
    respuesta = index.as_query_engine().query(query)
    
    # Respuesta de OpenAI basada en la consulta y agregarla al historial
    response = openai.ChatCompletion.create(
        model='gpt-3.5-turbo-0613',
        max_tokens=1000,
        messages=full_prompt,
        temperature=0.5,
        stream=True
    )

    partial_message = ""
    for chunk in response:
        if len(chunk['choices'][0]['delta']) != 0:
            partial_message = partial_message + chunk['choices'][0]['delta']['content']
            yield partial_message

# Define la interfaz de Gradio sin chat
gr.Interface(
    fn=predict,
    inputs=gr.Textbox("Escribe una pregunta...", label="Pregunta"),
    outputs=gr.Textbox(label="Respuesta"),
    title="Titulo titulo",
    description="Descripción descripción",
    theme="soft",
    examples=["¿Se puede tomar con anticonceptivos?", "Hola"],
    cache_examples=True,
    #live=True,
    button="Preguntar"
    clearbutton=None,
).launch()