app.py

import gradio as gr
import os
from datetime import datetime
from huggingface_hub import InferenceClient

class DefenseGenerator:
    def __init__(self):
        self.api_key = os.getenv("HF_API_KEY")
        if not self.api_key:
            raise EnvironmentError("HF_API_KEY não configurada")
        self.client = InferenceClient(api_key=self.api_key)
        
        self.document_types = {
            "resposta_acusacao": "Resposta à Acusação",
            "habeas_corpus": "Habeas Corpus",
            "revogacao_preventiva": "Revogação de Prisão Preventiva",
            "liberdade_provisoria": "Liberdade Provisória",
            "relaxamento_prisao": "Relaxamento de Prisão",
            "apelacao": "Recurso de Apelação",
            "embargos": "Embargos de Declaração"
        }

    def generate_content(self, prompt: str) -> str:
        try:
            response = self.client.chat.completions.create(
                model="mistralai/Mistral-7B-Instruct-v0.3",
                messages=[
                    {"role": "system", "content": "Você é um advogado criminalista experiente."},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=2000,
                temperature=0.7
            )
            return response.choices[0].message["content"]
        except Exception as e:
            return f"Erro na geração: {str(e)}"

    def gerar_documento(self, tipo: str, fatos: str, dados_basicos: dict) -> str:
        prompt = f"""Gere uma peça de {self.document_types[tipo]} completa com base nos seguintes fatos:

{fatos}

Inclua todos os elementos necessários:
- Qualificação completa das partes
- Fundamentos jurídicos pertinentes
- Jurisprudência relevante
- Doutrina aplicável
- Preliminares quando cabíveis
- Mérito detalhado
- Pedidos específicos

Dados básicos para formatação:
Comarca: {dados_basicos['comarca']}
Vara: {dados_basicos['vara']}
Cliente: {dados_basicos['cliente']}
OAB: {dados_basicos['oab']}"""

        return self.generate_content(prompt)

def interface():
    gerador = DefenseGenerator()
    
    with gr.Blocks() as demo:
        gr.Markdown("# Gerador de Peças Criminais")
        
        with gr.Row():
            tipo = gr.Dropdown(
                choices=list(gerador.document_types.keys()),
                label="Tipo de Peça",
                value="resposta_acusacao"
            )
        
        with gr.Row():
            with gr.Column():
                comarca = gr.Textbox(label="Comarca", placeholder="Ex: São Paulo")
                vara = gr.Textbox(label="Vara", placeholder="Ex: 1ª")
                
            with gr.Column():
                cliente = gr.Textbox(label="Nome do Cliente", placeholder="Ex: João da Silva")
                oab = gr.Textbox(label="OAB", placeholder="Ex: SP123456")
        
        fatos = gr.Textbox(
            label="Descreva os fatos e o que deseja na peça",
            placeholder="Descreva a situação do cliente e o que precisa ser argumentado...",
            lines=5
        )
        
        btn = gr.Button("Gerar")
        output = gr.Textbox(label="Peça Gerada", lines=20, show_copy_button=True)
        
        def process(tipo, fatos, comarca, vara, cliente, oab):
            dados = {
                "comarca": comarca,
                "vara": vara,
                "cliente": cliente,
                "oab": oab
            }
            return gerador.gerar_documento(tipo, fatos, dados)
            
        btn.click(
            fn=process,
            inputs=[tipo, fatos, comarca, vara, cliente, oab],
            outputs=output
        )
        
        exemplo = gr.Examples(
            examples=[
                ["resposta_acusacao", "Cliente preso em flagrante por furto de celular iPhone 13 da vítima Maria em 10/01/2024. É primário, tem residência fixa e trabalho. A vítima não reconheceu com certeza.", "São Paulo", "1ª", "João Silva", "SP123456"],
                ["habeas_corpus", "Cliente preso preventivamente há 90 dias por tráfico. Apreendidas 10g de maconha, é primário, tem residência e trabalho fixos. Não foram apreendidos outros materiais típicos de tráfico.", "São Paulo", "2ª", "Pedro Santos", "SP654321"]
            ],
            inputs=[tipo, fatos, comarca, vara, cliente, oab]
        )
    
    return demo

if __name__ == "__main__":
    demo = interface()
    demo.launch()