app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os
from datetime import datetime

DOCUMENT_TYPES = {
    "resposta_acusacao": "Resposta à Acusação (Art. 396-A do CPP)",
    "habeas_corpus": "Habeas Corpus (Art. 647 do CPP)", 
    "revogacao_preventiva": "Revogação de Prisão Preventiva (Art. 316 do CPP)",
    "liberdade_provisoria": "Liberdade Provisória (Art. 321 do CPP)",
    "relaxamento_prisao": "Relaxamento de Prisão (Art. 5º, LXV, CF)",
    "apelacao": "Recurso de Apelação (Art. 593 do CPP)",
    "embargos": "Embargos de Declaração (Art. 382 do CPP)"
}

TEMPLATE_PENAL = """EXCELENTÍSSIMO(A) SENHOR(A) DOUTOR(A) JUIZ(A) DE DIREITO DA {vara} VARA CRIMINAL DA COMARCA DE {comarca}/{estado}

AUTOS Nº {num_processo}

{tipo_peca}

I - QUALIFICAÇÃO DAS PARTES
{nome}, brasileiro(a), {idade} anos, {estado_civil}, {profissao}, RG nº {rg}, CPF nº {cpf}, residente e domiciliado à {endereco}, {cidade}/{estado}, vem, por seu advogado que esta subscreve, com fundamento na legislação aplicável, apresentar {tipo_documento}, pelos fatos e fundamentos a seguir expostos.

II - FUNDAMENTOS JURÍDICOS
1. {fundamentos}
2. {elementos_tipo_penal}

III - JURISPRUDÊNCIA APLICÁVEL
1. {jurisprudencia_1}
2. {jurisprudencia_2}

IV - DOUTRINA APLICÁVEL
1. {doutrina_1}
2. {doutrina_2}

V - PRELIMINARES
1. {preliminar_1}
2. {preliminar_2}

VI - MÉRITO
1. {merito_1}
2. {merito_2}
3. {merito_3}

VII - PEDIDOS
1. PRELIMINARMENTE:
   a) {pedido_preliminar_1}
   b) {pedido_preliminar_2}

2. NO MÉRITO:
   a) {pedido_merito_1}
   b) {pedido_merito_2}

3. PROVAS A PRODUZIR:
   {provas}

{cidade}, {data}

{advogado}
OAB/{estado} {oab}

ROL DE TESTEMUNHAS:
{testemunhas}
"""

class LegalDocGenerator:
    def __init__(self):
        self.api_key = os.getenv("HF_API_KEY")
        if not self.api_key:
            raise EnvironmentError("HF_API_KEY não configurada")
        self.client = InferenceClient(api_key=self.api_key)

    def _format_date(self) -> str:
        return datetime.now().strftime("%d de %B de %Y")

    def _build_prompt(self, doc_type: str, dados: dict) -> str:
        # Instruções base para o modelo
        system_prompt = """Você é um advogado criminalista especialista com anos de experiência.
Ao gerar peças processuais:
1. Use linguagem jurídica formal e técnica
2. Cite jurisprudência atual e relevante
3. Faça referências doutrinárias precisas
4. Desenvolva argumentos específicos para o caso
5. Mantenha coerência em toda a peça
6. Adapte os argumentos ao tipo de documento
7. Use legislação atualizada e pertinente"""

        # Prompt específico para o caso
        case_prompt = f"""TIPO DE PEÇA: {DOCUMENT_TYPES[doc_type]}

DADOS DO CASO:
Cliente: {dados['nome']}
Fatos: {dados['fatos']}
Comarca: {dados['comarca']}
Vara: {dados['vara']}

ELEMENTOS ESPECÍFICOS A INCLUIR:
1. Fundamentos específicos sobre {doc_type}
2. Jurisprudência favorável recente
3. Argumentos técnicos sobre {doc_type}
4. Pedidos adequados ao tipo de peça
5. Provas pertinentes ao caso

BASE PARA GERAÇÃO:
{TEMPLATE_PENAL}"""

        return system_prompt, case_prompt

    def generate(self, doc_type: str, dados: dict) -> str:
        try:
            system_prompt, case_prompt = self._build_prompt(doc_type, dados)
            dados['data'] = self._format_date()
            
            response = self.client.chat.completions.create(
                model="mistralai/Mistral-7B-Instruct-v0.3",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": case_prompt}
                ],
                max_tokens=2500,
                temperature=0.7
            )
            
            return response.choices[0].message["content"]
        except Exception as e:
            return f"Erro na geração: {str(e)}"

def create_interface():
    def process_submit(doc_type, fatos, comarca, vara, cliente, oab, proc_num, 
                      idade, estado_civil, profissao, rg, cpf, endereco, testemunhas):
        dados = {
            "tipo_peca": DOCUMENT_TYPES[doc_type],
            "nome": cliente,
            "fatos": fatos,
            "comarca": comarca,
            "vara": vara,
            "oab": oab,
            "num_processo": proc_num,
            "idade": idade,
            "estado_civil": estado_civil,
            "profissao": profissao,
            "rg": rg,
            "cpf": cpf,
            "endereco": endereco,
            "testemunhas": testemunhas
        }
        return LegalDocGenerator().generate(doc_type, dados)

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# Sistema de Geração de Peças Criminais v2.0")
        
        with gr.Tab("Dados do Processo"):
            with gr.Row():
                tipo = gr.Dropdown(choices=list(DOCUMENT_TYPES.keys()), 
                                 value="resposta_acusacao",
                                 label="Tipo de Peça")
                proc_num = gr.Textbox(label="Número do Processo",
                                    placeholder="0000000-00.0000.0.00.0000")
            
            with gr.Row():
                comarca = gr.Textbox(label="Comarca")
                vara = gr.Textbox(label="Vara Criminal")

        with gr.Tab("Dados do Cliente"):
            with gr.Row():
                cliente = gr.Textbox(label="Nome Completo")
                idade = gr.Number(label="Idade")
                estado_civil = gr.Dropdown(label="Estado Civil",
                                         choices=["Solteiro(a)", "Casado(a)", "Divorciado(a)", "Viúvo(a)"])
            
            with gr.Row():
                profissao = gr.Textbox(label="Profissão")
                rg = gr.Textbox(label="RG")
                cpf = gr.Textbox(label="CPF")
            
            endereco = gr.Textbox(label="Endereço Completo")

        with gr.Tab("Dados do Caso"):
            fatos = gr.TextArea(label="Descrição dos Fatos",
                              placeholder="Descreva detalhadamente a situação...",
                              lines=5)
            
            testemunhas = gr.TextArea(label="Rol de Testemunhas",
                                    placeholder="Nome completo e qualificação de cada testemunha",
                                    lines=3)

        with gr.Tab("Dados do Advogado"):
            oab = gr.Textbox(label="Número da OAB com Estado")

        doc_output = gr.TextArea(label="Documento Gerado",
                               lines=30,
                               show_copy_button=True)
        
        submit_btn = gr.Button("Gerar Peça Processual", variant="primary")
        submit_btn.click(
            fn=process_submit,
            inputs=[tipo, fatos, comarca, vara, cliente, oab, proc_num,
                   idade, estado_civil, profissao, rg, cpf, endereco, testemunhas],
            outputs=doc_output
        )

        gr.Examples([
            ["resposta_acusacao", 
             "Cliente preso em flagrante por furto de celular iPhone 13. É primário, tem residência fixa e trabalho lícito comprovado. A vítima não o reconheceu com certeza.",
             "São Paulo", "1ª", "João Silva", "SP123456", "0001234-56.2024.8.26.0000",
             35, "Casado(a)", "Comerciante", "12.345.678-9", "123.456.789-00",
             "Rua das Flores, 123, Centro", 
             "1. Maria Oliveira - Testemunha ocular\n2. José Santos - Empregador do acusado"]
        ], inputs=[tipo, fatos, comarca, vara, cliente, oab, proc_num,
                  idade, estado_civil, profissao, rg, cpf, endereco, testemunhas])

    return demo

if __name__ == "__main__":
    create_interface().launch()