from openai import OpenAI
import gradio as gr

# Configuración del cliente OpenAI con la API de NVIDIA
client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key="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"
)

# Instrucciones para el modelo
instrucciones = """
Eres un Ministro de la Corte Suprema de Chile, especializado en derecho civil. Tu tarea principal es mejorar borradores de resoluciones judiciales y responder preguntas generales, asegurando un tono formal y técnico cuando sea necesario, y proporcionando respuestas claras y precisas. Tus habilidades incluyen:
1. Utilizar un tono formal y técnico para garantizar claridad y coherencia.
2. Usar terminología legal específica, como 'folio', 'autos', 'resuelve', 'artículo', 'Código de Procedimiento Civil', 'audiencia', 'notificación', 'prueba', 'nulidad', 'oficiar' y 'notificar'.
3. Crear oraciones complejas y subordinadas, manteniendo un tono formal e imperativo.
4. Proporcionar referencias detalladas a regulaciones y precedentes judiciales, citando consistentemente artículos de leyes y sentencias anteriores.
5. Mantener un tono uniforme y detallado en todo el documento, centrado en la precisión terminológica.
"""

instrucciones_legales = """
Revisas el documento para asegurar:
- Comprensión del contexto y propósito.
- Identificación del tipo de documento judicial (sentencia, apelación).
- Progresión lógica de ideas y coherencia entre secciones.
- Claridad y precisión en cada párrafo y oración.
- Uso consistente de terminología legal.
- Corrección gramatical y claridad en la estructura de las oraciones.
- Uso adecuado de conectores para transiciones suaves.
- Precisión en términos legales y sugerencias de expresiones más efectivas.
- Consideración de sugerencias y validaciones, clasificándolas según su relevancia.
- Revisión integral y presentación concisa y unificada de la información.
"""

def mejorar_resolucion(input_text):
    # Construcción del prompt con instrucciones unificadas
    prompt = f"Adopta el rol de Ministro de la Corte Suprema de Chile y responde a la siguiente pregunta/pregunta con base en las instrucciones e instrucciones legales proporcionadas:\n{instrucciones}\n{instrucciones_legales}\nTexto o pregunta:\n{input_text}\nRespuesta:"

    try:
        completion = client.chat.completions.create(
            model="meta/llama-3.1-405b-instruct",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.5,
            top_p=0.7,
            max_tokens=1024,
            stream=True
        )

        response_content = ""
        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                response_content += chunk.choices[0].delta.content
        return response_content.strip()

    except Exception as e:
        return f"Error: {e}"

# Definición de la interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Mejora de Resoluciones Judiciales y Respuestas con Nemotron")
    input_text = gr.Textbox(label="Introduce tu resolución judicial o pregunta")
    output_text = gr.Textbox(label="Respuesta mejorada o respuesta corta")
    submit_button = gr.Button("Enviar")
    submit_button.click(fn=mejorar_resolucion, inputs=input_text, outputs=output_text)

# Lanzamiento de la aplicación
demo.launch()