app.py

import gradio as gr
from transformers import pipeline
import json
from datetime import datetime
import logging
from typing import Dict, List, Optional

# Configuração do logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class PromptManager:
    """Gerencia a criação e otimização de prompts"""
    
    @staticmethod
    def get_base_prompt(doc_type: str, context: Dict[str, str]) -> str:
        """Retorna o prompt base para cada tipo de documento"""
        
        prompts = {
            "habeas_corpus": f"""
            Atue como um advogado criminalista experiente e gere um Habeas Corpus detalhado:
            
            DADOS DO CASO:
            Cliente: {context.get('client_name')}
            Processo: {context.get('process_number')}
            Tribunal: {context.get('court')}
            Comarca: {context.get('jurisdiction')}
            
            FATOS:
            {context.get('facts')}
            
            FUNDAMENTOS JURÍDICOS:
            {context.get('legal_basis')}
            
            Gere um Habeas Corpus completo e tecnicamente preciso, incluindo:
            1. Cabeçalho com endereçamento correto
            2. Qualificação das partes
            3. Fatos relevantes
            4. Fundamentos jurídicos
            5. Jurisprudência relevante
            6. Pedidos (liminar e mérito)
            """,
            
            "denuncia": f"""
            Atue como um promotor de justiça experiente e gere uma denúncia criminal:
            
            DADOS DO CASO:
            Processo: {context.get('process_number')}
            Tribunal: {context.get('court')}
            Comarca: {context.get('jurisdiction')}
            
            FATOS:
            {context.get('facts')}
            
            ENQUADRAMENTO LEGAL:
            {context.get('legal_basis')}
            
            Gere uma denúncia criminal completa, incluindo:
            1. Qualificação do(s) denunciado(s)
            2. Narrativa clara dos fatos
            3. Tipificação penal
            4. Circunstâncias do crime
            5. Rol de testemunhas
            6. Pedidos
            """,
            
            "alegacoes_finais": f"""
            Atue como um advogado criminalista e gere alegações finais:
            
            DADOS DO PROCESSO:
            Cliente: {context.get('client_name')}
            Processo: {context.get('process_number')}
            Tribunal: {context.get('court')}
            
            FATOS E PROVAS:
            {context.get('facts')}
            
            TESE DEFENSIVA:
            {context.get('legal_basis')}
            
            Elabore alegações finais completas, incluindo:
            1. Relatório do processo
            2. Preliminares (se houver)
            3. Mérito
            4. Análise das provas
            5. Pedidos
            """
        }
        
        return prompts.get(doc_type, "Tipo de documento não suportado")

class DocumentGenerator:
    """Gerencia a geração de documentos usando GPT4All-J"""
    
    def __init__(self):
        self.pipe = pipeline("text-generation", model="nomic-ai/gpt4all-j")
        self.prompt_manager = PromptManager()
    
    def generate(self, doc_type: str, context: Dict[str, str]) -> str:
        """Gera o documento baseado no tipo e contexto"""
        try:
            prompt = self.prompt_manager.get_base_prompt(doc_type, context)
            
            # Configuração para a geração
            result = self.pipe(
                prompt,
                max_length=2048,
                temperature=0.7,
                top_p=0.95,
                num_return_sequences=1,
                do_sample=True
            )[0]['generated_text']
            
            return self._format_output(result)
            
        except Exception as e:
            logger.error(f"Erro na geração: {str(e)}")
            return f"Erro na geração do documento: {str(e)}"
    
    def _format_output(self, text: str) -> str:
        """Formata o texto gerado para melhor apresentação"""
        return text.strip()

class WebInterface:
    """Interface Gradio para o gerador de documentos"""
    
    def __init__(self):
        self.generator = DocumentGenerator()
        self.create_interface()
    
    def create_interface(self):
        """Cria a interface web com Gradio"""
        
        with gr.Blocks(theme=gr.themes.Soft()) as self.app:
            gr.Markdown("""
            # Gerador de Documentos Criminais
            ### Powered by GPT4All-J
            """)
            
            with gr.Row():
                with gr.Column():
                    # Seleção do tipo de documento
                    doc_type = gr.Dropdown(
                        choices=[
                            "Habeas Corpus",
                            "Denúncia",
                            "Alegações Finais"
                        ],
                        label="Tipo de Documento",
                        value="Habeas Corpus"
                    )
                    
                    # Informações do processo
                    with gr.Group():
                        gr.Markdown("### Informações do Processo")
                        client_name = gr.Textbox(
                            label="Nome do Cliente",
                            placeholder="Nome completo do cliente"
                        )
                        process_number = gr.Textbox(
                            label="Número do Processo",
                            placeholder="NNNNNNN-NN.NNNN.N.NN.NNNN"
                        )
                        court = gr.Textbox(
                            label="Tribunal",
                            value="TRIBUNAL DE JUSTIÇA DO ESTADO"
                        )
                        jurisdiction = gr.Textbox(
                            label="Comarca",
                            placeholder="Comarca onde tramita o processo"
                        )
                    
                    # Detalhes do caso
                    with gr.Group():
                        gr.Markdown("### Detalhes do Caso")
                        facts = gr.Textbox(
                            label="Fatos",
                            lines=5,
                            placeholder="Descreva os fatos relevantes do caso..."
                        )
                        legal_basis = gr.Textbox(
                            label="Fundamentos Jurídicos",
                            lines=3,
                            placeholder="Indique os fundamentos legais principais..."
                        )
                    
                    # Status e botões
                    status = gr.Textbox(
                        label="Status",
                        interactive=False
                    )
                    
                    generate_btn = gr.Button(
                        "Gerar Documento",
                        variant="primary"
                    )
                
                with gr.Column():
                    output = gr.Textbox(
                        label="Documento Gerado",
                        lines=30,
                        show_copy_button=True
                    )
            
            # Exemplos de uso
            gr.Examples(
                examples=[
                    [
                        "Habeas Corpus",
                        "João da Silva",
                        "0000123-45.2024.8.26.0000",
                        "TRIBUNAL DE JUSTIÇA DO ESTADO DE SÃO PAULO",
                        "São Paulo",
                        "Paciente preso em flagrante no dia 25/12/2024 por suposto furto...",
                        "Art. 5º, LXVIII da CF/88; Art. 647 do CPP; Ausência dos requisitos do Art. 312 do CPP"
                    ]
                ],
                inputs=[
                    doc_type, client_name, process_number,
                    court, jurisdiction, facts, legal_basis
                ]
            )
            
            # Eventos
            generate_btn.click(
                fn=self._generate_document,
                inputs=[
                    doc_type, client_name, process_number,
                    court, jurisdiction, facts, legal_basis
                ],
                outputs=[output, status]
            )
    
    def _generate_document(
        self, doc_type: str, client_name: str,
        process_number: str, court: str,
        jurisdiction: str, facts: str,
        legal_basis: str
    ) -> tuple:
        """Gera o documento com os parâmetros fornecidos"""
        
        try:
            # Validação básica
            if not all([client_name, process_number, facts, legal_basis]):
                return "Erro: Todos os campos obrigatórios devem ser preenchidos", "⚠️ Preencha todos os campos obrigatórios"
            
            context = {
                "client_name": client_name,
                "process_number": process_number,
                "court": court,
                "jurisdiction": jurisdiction,
                "facts": facts,
                "legal_basis": legal_basis
            }
            
            # Atualiza status
            yield "", "⏳ Gerando documento..."
            
            # Gera documento
            result = self.generator.generate(
                doc_type.lower().replace(" ", "_"),
                context
            )
            
            return result, "✅ Documento gerado com sucesso!"
            
        except Exception as e:
            logger.error(f"Erro: {str(e)}")
            return "", f"❌ Erro: {str(e)}"
    
    def launch(self):
        """Inicia a interface web"""
        self.app.launch(share=True)

if __name__ == "__main__":
    # Dependências necessárias:
    # pip install gradio transformers torch
    
    interface = WebInterface()
    interface.launch()