app.py

import gradio as gr
from huggingface_hub import InferenceClient
from datetime import datetime
import logging
from typing import Dict, List, Optional

# Configuração do logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class DocumentGenerator:
    """Gerencia a geração de documentos usando Mistral via InferenceClient"""
    
    def __init__(self):
        # Inicializa o cliente de inferência
        self.client = InferenceClient()
        self.model = "mistralai/Mistral-7B-Instruct-v0.2"
    
    def generate(self, doc_type: str, context: Dict[str, str]) -> str:
        """Gera o documento baseado no tipo e contexto"""
        try:
            # Prepara a mensagem para o modelo
            system_prompt = """Você é um advogado criminalista experiente especializado em gerar documentos jurídicos no Brasil.
            Gere o documento solicitado usando linguagem formal, técnica e no formato jurídico correto."""
            
            user_content = f"""Gere um {doc_type} com os seguintes dados:
            
            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')}
            
            O documento deve ser completo, tecnicamente preciso e seguir o formato jurídico brasileiro."""

            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_content}
            ]

            # Faz a chamada ao modelo
            completion = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                max_tokens=2048,
                temperature=0.3,
                top_p=0.85,
            )

            return self._format_output(completion.choices[0].message.content)
            
        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"""
        if not text:
            return "Erro: Nenhum texto gerado"
        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 Mistral-7B-Instruct
            """)
            
            with gr.Row():
                with gr.Column():
                    # Seleção do tipo de documento
                    doc_type = gr.Dropdown(
                        choices=[
                            "Habeas Corpus",
                            "Denúncia",
                            "Alegações Finais",
                            "Resposta à Acusação",
                            "Recurso em Sentido Estrito",
                            "Apelação Criminal"
                        ],
                        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..."
                        )
                    
                    generate_btn = gr.Button(
                        "Gerar Documento",
                        variant="primary"
                    )
                
                with gr.Column():
                    output = gr.Textbox(
                        label="Documento Gerado",
                        lines=30,
                        show_copy_button=True
                    )
                    
                    status = gr.Textbox(
                        label="Status",
                        interactive=False
                    )
            
            # Exemplo 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,
                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()

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