app.py

import gradio as gr
from gradio_client import Client
import json
from datetime import datetime
import os
from typing import Dict, List, Optional
import logging

# 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 a API do LLama"""
    
    def __init__(self):
        self.client = Client("ysharma/Chat_with_Meta_llama3_8b")
        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)
            
            result = self.client.predict(
                message=prompt,
                request=0.85,  # temperatura mais baixa para maior consistência
                param_3=3072,  # tokens aumentados para documentos longos
                api_name="/chat"
            )
            
            return self._format_output(result)
            
        except Exception as e:
            logger.error(f"Erro na geração do documento: {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 LLM Technology
            """)
            
            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
                    )
            
            # 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
            )
    
    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"""
        
        # Validação básica
        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.lower().replace(" ", "_"),
            context
        )
    
    def launch(self):
        """Inicia a interface web"""
        self.app.launch(share=True)

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