app.py

import os
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from dotenv import load_dotenv

# Cargar variables de entorno
# load_dotenv()

# Sistema de prompt mejorado (igual que antes)
SYSTEM_PROMPT = """Responde con precisión sobre el delito que se consulta, citando la legislación vigente en España y proporcionando contexto legal actualizado:
   - Menciona siempre los artículos específicos del Código Penal español al inicio.
   - Si el delito ha sido modificado en reformas recientes, indica en qué consistió la reforma.
   - Define el **tipo básico del delito**.
   - Explica los **subtipos agravados y atenuados** si existen.
   - Cita los artículos correspondientes.
   - Menciona sentencias recientes del Tribunal Supremo cuando sea posible."""

# Cargar modelo y tokenizador de Hugging Face
MODEL_NAME = "EleutherAI/gpt-neo-2.7B"  # Puedes cambiarlo por otro modelo (ej: "tiiuae/falcon-7b" -ram exceeded- o  "meta-llama/Llama-2-7b-chat-hf" -se debe solictar acesso en meta ai-)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# Crear pipeline de generación de texto
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Historial global de chat
chat_history = []

def consulta_penalista(mensaje, tipo_consulta):
    """Procesa una consulta jurídica y devuelve la respuesta del modelo."""
    global chat_history
    
    # Crear un prompt que indique directamente qué responder
    prompt = f"""A continuación hay una consulta jurídica sobre el sistema penal español.

Consulta: {mensaje}

Respuesta sobre el delito consultado según la legislación española:
- Artículos del Código Penal español: """
    
    try:
        # Crear el prompt completo (sistema + historial)
        # prompt = SYSTEM_PROMPT + "\n\n" + "\n".join([msg["content"] for msg in chat_history])
        prompt = SYSTEM_PROMPT + mensaje
        
        # Generar respuesta con el modelo
        respuesta = generator(
            prompt,
            max_length=1000,  # Longitud máxima de la respuesta (antes 500)
            num_return_sequences=1,  # Número de respuestas a generar

            # (NUEVO)
            temperature=0.5,  # Controlar la creatividad (valores más bajos = más determinista)
            top_p=0.9,  # Usar muestreo de núcleo (nucleus sampling)
            truncation=True,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=True
        )
        
        respuesta_texto = respuesta[0]["generated_text"]
        
        # Añadir la respuesta del asistente al historial
        chat_history.append({"role": "assistant", "content": respuesta_texto})
        
        # Formatear el historial para mostrarlo en la interfaz
        historial_formateado = "\n\n".join([
            f"{'Usuario' if msg['role'] == 'user' else 'Asistente'}: {msg['content']}"
            for msg in chat_history
        ])
        
        return respuesta_texto, historial_formateado
    
    except Exception as e:
        return f"Error al procesar la solicitud: {str(e)}", "\n".join(str(m) for m in chat_history)

def limpiar_historial():
    """Limpia el historial de conversación."""
    global chat_history
    chat_history = []
    return "", "Historial limpiado."

# Interfaz de usuario con Gradio (igual que antes)
interfaz = gr.Blocks()

with interfaz:
    gr.Markdown("# ⚖️ Asistente Jurídico Penal")
    gr.Markdown("Realiza consultas sobre derecho penal y recibe respuestas fundamentadas.")
    
    tipo_consulta = gr.Radio(
        choices=["Asesoría General", "Jurisprudencia", "Redacción de Documentos"],
        label="Tipo de Consulta",
        value="Asesoría General"
    )
    
    entrada_texto = gr.Textbox(
        label="Pregunta Jurídica",
        placeholder="Escribe tu consulta aquí..."
    )
    
    with gr.Row():
        boton_enviar = gr.Button("Consultar")
        boton_limpiar = gr.Button("Limpiar Historial")
    
    salida_texto = gr.Textbox(label="Respuesta")
    historial_texto = gr.Textbox(
        label="Historial de Conversación",
        interactive=False,
        lines=10
    )
    
    boton_enviar.click(
        consulta_penalista,
        inputs=[entrada_texto, tipo_consulta],
        outputs=[salida_texto, historial_texto]
    )
    
    boton_limpiar.click(
        limpiar_historial,
        outputs=[salida_texto, historial_texto]
    )

interfaz.launch()