import os
import openai
from llama_index import StorageContext, load_index_from_storage, LLMPredictor, ServiceContext
from llama_index.tools import QueryEngineTool, ToolMetadata
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.")




#storage_context

# reconstruir storage context Doc_1
storage_context = StorageContext.from_defaults(persist_dir=".")
modelo = LLMPredictor(llm=ChatOpenAI(temperature=0, model_name='gpt-3.5-turbo-16k', max_tokens=7000,top_p=0))
service_context = ServiceContext.from_defaults(llm_predictor=modelo)


# Cargar index Doc_1
index = load_index_from_storage(service_context=service_context, storage_context=storage_context)


#Query Engine
query_engine = index.as_query_engine()
query_engine_tools = [
    QueryEngineTool(
        query_engine=query_engine,
        metadata=ToolMetadata(
             name='Prospecto medicamento Acsodix con instrucciones para su uso y precauciones',
            description='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. Sobredosis: Síntomas y tratamientos en caso de sobredosis. Presentación: Cajas conteniendo 30 comprimidos recubiertos en diferentes dosis. Almacenamiento: Conservar a temperatura ambiente inferior a 30 °C. Fabricante y Distribuidor: Mega Labs S.A., Industria Uruguaya, 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.')
    )
]
engine = index.as_query_engine(query_engine_tools=query_engine_tools, service_context=service_context,streaming=True,similarity_top_k=5)






# Prompt personalizado para el nuevo index
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. Sobredosis: Síntomas y tratamientos en caso de sobredosis. Presentación: Cajas conteniendo 30 comprimidos recubiertos en diferentes dosis. Almacenamiento: Conservar a temperatura ambiente inferior a 30 °C. Fabricante y Distribuidor: Mega Labs S.A., Industria Uruguaya, 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.'

# Funcion de chat para el nuevo index Doc_1
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"]
    
    respuesta = engine.query(user_message)
    
    response = openai.ChatCompletion.create(
        model='gpt-3.5-turbo-16k',  # Use the new model
        max_tokens=1000,
        messages=full_prompt,
        temperature=0.5,
        stream=True
    )

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

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