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 OpenHermes"""
    
    def __init__(self):
        self.pipe = pipeline(
            "text-generation",
            model="jondurbin/airoboros-l2-70b-gpt4-1.4.1",  # Modelo público com bom desempenho em PT
            trust_remote_code=True
        )
        self.prompt_manager = PromptManager()
    
    def generate(self, doc_type: str, context: Dict[str, str]) -> str:
        """Gera o documento baseado no tipo e contexto"""
        try:
            # Prepara o prompt com instruções específicas
            base_prompt = self.prompt_manager.get_base_prompt(doc_type, context)
            
            system_prompt = """Você é um advogado criminalista experiente especializado em gerar documentos jurídicos.
            Gere o documento solicitado usando linguagem formal e técnica apropriada."""
            
            formatted_prompt = f"""### System:
            {system_prompt}
            
            ### User:
            {base_prompt}
            
            ### Assistant:
            Gerando o documento solicitado em formato jurídico adequado:
            """
            
            # Configuração otimizada para documentos jurídicos
            result = self.pipe(
                formatted_prompt,
                max_new_tokens=2048,          # Aumentado para documentos longos
                temperature=0.3,              # Reduzido para maior consistência
                top_p=0.85,                   # Ajustado para melhor coerência
                top_k=40,                     # Limita as opções de tokens
                num_return_sequences=1,
                do_sample=True,
                repetition_penalty=1.2,       # Evita repetições
                pad_token_id=self.pipe.tokenizer.eos_token_id
            )[0]['generated_text']
            
            # Remove o prompt do resultado e retorna apenas o documento
            final_text = result.split("### Assistant:")[-1].strip()
            return self._format_output(final_text)
            
        except Exception as e:
            logger.error(f"Erro na geração: {str(e)}")
            if "CUDA" in str(e):
                return "Erro: Problema com recursos GPU. Tentando novamente com configurações reduzidas..."
            elif "out of memory" in str(e):
                return "Erro: Memória insuficiente. Tente um documento menor."
            else:
                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"""
        if not isinstance(text, str):
            if isinstance(text, (list, tuple)) and len(text) > 0:
                text = str(text[0])
            else:
                text = str(text)
                
        # Remove linhas vazias extras
        lines = [line for line in text.split('\n') if line.strip()]
        formatted_text = '\n'.join(lines)
        
        # Ajusta espaçamento de parágrafos
        formatted_text = formatted_text.replace('\n\n\n', '\n\n')
        
        return formatted_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()