app.py

import os
import gradio as gr
from huggingface_hub import HfApi, InferenceApi
from datetime import datetime
import logging
from typing import Dict, List

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

class DocumentGenerator:
    """Gerencia a geração de documentos usando modelo público"""
    
    def __init__(self):
        # Usando um Space público que está efetivamente disponível
        self.client = InferenceApi(
            repo_id="OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
            token=os.environ.get("HF_TOKEN")
        )
    
    def generate(self, doc_type: str, context: Dict[str, str]) -> str:
        """Gera o documento usando o modelo"""
        try:
            prompt = f"""You are a Brazilian criminal lawyer. Create a {doc_type} in Portuguese following Brazilian legal standards.

Information:
Client: {context.get('client_name')}
Process: {context.get('process_number')}
Court: {context.get('court')}
Jurisdiction: {context.get('jurisdiction')}

Facts:
{context.get('facts')}

Legal Basis:
{context.get('legal_basis')}

Instructions:
1. Use formal legal Portuguese
2. Follow Brazilian legal document format
3. Include all required sections
4. Be precise and clear
5. Keep proper legal formatting"""

            response = self.client(
                inputs=prompt,
                parameters={
                    "max_new_tokens": 2048,
                    "temperature": 0.3,
                    "top_p": 0.95,
                    "repetition_penalty": 1.15,
                    "do_sample": True
                }
            )

            return self._format_output(response[0]["generated_text"])
            
        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 o prompt da resposta
        text = text.split("Instructions:")[-1].strip()
        
        # Ajusta formatação
        lines = [line.strip() for line in text.split('\n') if line.strip()]
        formatted_text = '\n\n'.join(lines)
        
        # Adiciona data atual
        current_date = datetime.now().strftime('%d de %B de %Y')
        formatted_text = f"{formatted_text}\n\n{context.get('jurisdiction')}, {current_date}"
        
        return formatted_text

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("""
            # Criminal.ai - Gerador de Peças Processuais
            ### Sistema Inteligente para Geração de Documentos Jurídicos
            """)
            
            with gr.Row():
                with gr.Column():
                    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"
                    )
                    
                    with gr.Group():
                        gr.Markdown("### Informações do Processo")
                        client_name = gr.Textbox(
                            label="Nome do Cliente",
                            placeholder="Nome completo"
                        )
                        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="Nome da comarca"
                        )
                    
                    with gr.Group():
                        gr.Markdown("### Detalhes do Caso")
                        facts = gr.Textbox(
                            label="Fatos",
                            lines=5,
                            placeholder="Descreva os fatos relevantes..."
                        )
                        legal_basis = gr.Textbox(
                            label="Fundamentos Jurídicos",
                            lines=3,
                            placeholder="Fundamentos legais..."
                        )
                    
                    generate_btn = gr.Button("📝 Gerar Documento", variant="primary")
                
                with gr.Column():
                    output = gr.Textbox(
                        label="Documento Gerado",
                        lines=30,
                        show_copy_button=True
                    )
            
            # 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, CP). O paciente é primário, possui residência fixa no distrito da culpa e ocupação lícita comprovada. Não estão presentes os requisitos autorizadores da prisão preventiva.",
                        "Art. 5º, LXVIII, CF/88; Art. 647 do CPP; Ausência dos requisitos do Art. 312 do CPP. Súmula 308 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
            )
    
    def _generate_document(
        self, doc_type: str, client_name: str,
        process_number: str, court: str,
        jurisdiction: str, facts: str,
        legal_basis: str
    ) -> str:
        """Gera o documento com os parâmetros fornecidos"""
        
        try:
            if not all([client_name, process_number, facts, legal_basis]):
                return "Erro: Todos os campos obrigatórios devem ser preenchidos"
            
            context = {
                "client_name": client_name,
                "process_number": process_number,
                "court": court,
                "jurisdiction": jurisdiction,
                "facts": facts,
                "legal_basis": legal_basis
            }
            
            return self.generator.generate(doc_type, context)
            
        except Exception as e:
            logger.error(f"Erro: {str(e)}")
            return f"Erro na geração: {str(e)}"
    
    def launch(self):
        """Inicia a interface web"""
        self.app.launch(share=True)

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