app.py

import os
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__)

# Substitua com sua chave de API
HF_API_KEY = "hf_seu_token_aqui"

class DocumentManager:
    """Gerencia templates e contextos dos documentos"""

    @staticmethod
    def create_prompt_messages(doc_type: str, context: Dict[str, str]) -> List[Dict[str, str]]:
        """Cria a sequência de mensagens para o modelo"""
        
        system_message = {
            "role": "system",
            "content": """Você é um advogado criminalista brasileiro altamente experiente, especializado em redigir peças processuais.
            Gere documentos jurídicos formais, tecnicamente precisos e no formato do direito brasileiro."""
        }
        
        user_message = {
            "role": "user",
            "content": f"""Por favor, gere um {doc_type} completo e tecnicamente preciso com os seguintes dados:

QUALIFICAÇÃO:
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 seguir a formatação jurídica padrão, incluindo:
1. Cabeçalho com endereçamento correto
2. Qualificação completa das partes
3. Exposição clara dos fatos
4. Fundamentação jurídica sólida
5. Pedidos específicos
6. Fechamento formal com local, data e espaço para assinatura"""
        }
        
        return [system_message, user_message]

class DocumentGenerator:
    """Gerencia a geração de documentos usando HuggingFace Inference API"""
    
    def __init__(self, api_key: str):
        self.client = InferenceClient(api_key=api_key)
        self.model = "mistralai/Mistral-7B-Instruct-v0.2"
        self.doc_manager = DocumentManager()
    
    def generate(self, doc_type: str, context: Dict[str, str]) -> str:
        """Gera o documento usando o modelo"""
        try:
            # Prepara as mensagens
            messages = self.doc_manager.create_prompt_messages(doc_type, context)
            
            # Faz a chamada à API
            completion = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.3,    # Mais conservador para documentos formais
                top_p=0.85,         # Mantém coerência
                max_tokens=2048,    # Documento completo
            )
            
            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"
        
        # Remove espaços extras e formata parágrafos
        lines = [line.strip() for line in text.split('\n') if line.strip()]
        formatted_text = '\n\n'.join(lines)
        
        return formatted_text.strip()

class WebInterface:
    """Interface Gradio para o gerador de documentos"""
    
    def __init__(self):
        self.generator = DocumentGenerator(api_key=HF_API_KEY)
        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("""
            # Criminal.ai - Gerador de Peças Processuais
            ### Sistema Inteligente de Geração de Documentos Jurídicos
            """)
            
            with gr.Row():
                with gr.Column():
                    # Tipo de documento
                    doc_type = gr.Dropdown(
                        choices=[
                            "Habeas Corpus",
                            "Denúncia Criminal",
                            "Alegações Finais",
                            "Resposta à Acusação",
                            "Recurso em Sentido Estrito",
                            "Apelação Criminal"
                        ],
                        label="Tipo de Documento",
                        value="Habeas Corpus",
                        info="Selecione o tipo de peça processual"
                    )
                    
                    # 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/paciente",
                            info="Digite o nome completo"
                        )
                        
                        process_number = gr.Textbox(
                            label="Número do Processo",
                            placeholder="NNNNNNN-NN.NNNN.N.NN.NNNN",
                            info="Número CNJ do processo"
                        )
                        
                        court = gr.Textbox(
                            label="Tribunal",
                            value="TRIBUNAL DE JUSTIÇA DO ESTADO",
                            info="Nome completo do tribunal"
                        )
                        
                        jurisdiction = gr.Textbox(
                            label="Comarca",
                            placeholder="Nome da comarca",
                            info="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...",
                            info="Descreva os fatos de forma clara e objetiva"
                        )
                        
                        legal_basis = gr.Textbox(
                            label="Fundamentos Jurídicos",
                            lines=3,
                            placeholder="Indique os fundamentos legais...",
                            info="Artigos, jurisprudência e doutrina aplicáveis"
                        )
                    
                    # Status e botões
                    with gr.Row():
                        generate_btn = gr.Button(
                            "📝 Gerar Documento",
                            variant="primary",
                            scale=2
                        )
                        
                        clear_btn = gr.Button(
                            "🗑️ Limpar",
                            variant="secondary",
                            scale=1
                        )
                    
                    status = gr.Textbox(
                        label="Status",
                        interactive=False,
                        show_label=False
                    )
                
                # Coluna de saída
                with gr.Column():
                    output = gr.Textbox(
                        label="Documento Gerado",
                        lines=30,
                        show_copy_button=True,
                        show_label=True,
                        container=True
                    )
                    
                    with gr.Row():
                        copy_btn = gr.Button("📋 Copiar")
                        download_btn = gr.Button("💾 Salvar como .txt")
            
            # Exemplos
            gr.Examples(
                examples=[
                    [
                        "Habeas Corpus",
                        "João da Silva Santos",
                        "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 simples (art. 155, caput, do CP). Não estão presentes os requisitos da prisão preventiva, sendo o paciente primário, com residência fixa e trabalho lícito.",
                        "Art. 5º, LXVIII da CF/88; Art. 647 do CPP; Art. 312 do CPP; Súmula 308 do STJ."
                    ]
                ],
                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 clear_inputs():
                return [
                    gr.update(value="") for _ in range(6)
                ] + [gr.update(value="Campos limpos")]
            
            clear_btn.click(
                fn=clear_inputs,
                inputs=[],
                outputs=[
                    client_name, process_number, court,
                    jurisdiction, facts, legal_basis,
                    status
                ]
            )
            
            copy_btn.click(
                fn=lambda x: gr.update(value=x),
                inputs=[output],
                outputs=[output]
            )
            
            def save_document(text):
                if not text:
                    return "Nenhum documento para salvar"
                    
                filename = f"documento_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"
                with open(filename, 'w', encoding='utf-8') as f:
                    f.write(text)
                return f"Documento salvo como {filename}"
            
            download_btn.click(
                fn=save_document,
                inputs=[output],
                outputs=[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(share=True)

if __name__ == "__main__":
    # Verificar token
    if HF_API_KEY == "hf_seu_token_aqui":
        print("⚠️ Atenção: Configure sua API key do Hugging Face antes de executar!")
        print("Substitua 'hf_seu_token_aqui' pela sua chave real.")
    else:
        interface = WebInterface()
        interface.launch()