app.py

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

DOCUMENT_TYPES = {
    "resposta_acusacao": "Resposta à Acusação",
    "habeas_corpus": "Habeas Corpus", 
    "revogacao_preventiva": "Revogação de Prisão Preventiva",
    "liberdade_provisoria": "Liberdade Provisória",
    "relaxamento_prisao": "Relaxamento de Prisão",
    "apelacao": "Recurso de Apelação",
    "embargos": "Embargos de Declaração"
}

DEFAULT_SECTIONS = """AUTOS Nº [NUMERO_PROCESSO]

I - QUALIFICAÇÃO DAS PARTES
[Detalhes do acusado e seu advogado]

II - FUNDAMENTOS JURÍDICOS
1. [Base legal específica do caso]
2. [Elementos necessários para configuração do crime]

III - JURISPRUDÊNCIA APLICÁVEL
1. [Precedentes favoráveis]
2. [Entendimentos dos tribunais]

IV - DOUTRINA APLICÁVEL
1. [Citações doutrinárias relevantes]
2. [Posicionamentos doutrinários]

V - PRELIMINARES
1. [Questões processuais]
2. [Nulidades e vícios]

VI - MÉRITO
1. [Argumentos principais]
2. [Provas e contraprovas]

VII - PEDIDOS
1. [Pedido principal]
2. [Pedidos subsidiários]
3. [Provas a produzir]

[CIDADE], [DATA]

[ADVOGADO]
OAB/[ESTADO] [NÚMERO]

ROL DE TESTEMUNHAS:
1. [Nome e qualificação]
2. [Nome e qualificação]
3. [Nome e qualificação]"""

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 _build_prompt(self, doc_type: str, caso: dict) -> str:
        return f"""Atue como um advogado criminalista experiente gerando uma peça de {DOCUMENT_TYPES[doc_type]}.

CASO CONCRETO:
- Acusado: {caso['cliente']}
- Fatos: {caso['fatos']}
- Comarca: {caso['comarca']}
- Vara: {caso['vara']} Vara Criminal
- OAB: {caso['oab']}

INSTRUÇÕES ESPECÍFICAS:
1. Use linguagem jurídica formal
2. Cite jurisprudência real e atualizada
3. Fundamente com doutrina pertinente
4. Destaque elementos favoráveis ao acusado
5. Mantenha estrutura com numeração romana para seções principais
6. Inclua elementos de prova mencionados nos fatos
7. Desenvolva argumentos técnicos para cada pedido
8. Mantenha formatação profissional
9. Use artigos corretos do Código Penal e CPP

ESTRUTURA BASE:
{DEFAULT_SECTIONS}"""

    def generate(self, doc_type: str, caso: dict) -> str:
        try:
            prompt = self._build_prompt(doc_type, caso)
            response = self.client.chat.completions.create(
                model="mistralai/Mistral-7B-Instruct-v0.3",
                messages=[
                    {"role": "system", "content": "Você é um advogado criminalista especialista em gerar peças jurídicas técnicas e bem fundamentadas."},
                    {"role": "user", "content": 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(tipo, fatos, comarca, vara, cliente, oab, num_processo="", testemunhas=""):
        caso = {
            "cliente": cliente,
            "fatos": fatos,
            "comarca": comarca,
            "vara": vara,
            "oab": oab,
            "num_processo": num_processo,
            "testemunhas": testemunhas
        }
        return LegalDocGenerator().generate(tipo, caso)

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# Sistema Criminal.AI - Gerador de Peças Jurídicas")
        
        with gr.Row():
            tipo = gr.Dropdown(
                choices=list(DOCUMENT_TYPES.keys()),
                value="resposta_acusacao",
                label="Tipo de Peça"
            )
            num_processo = gr.Textbox(
                label="Número do Processo",
                placeholder="Ex: 0000000-00.0000.0.00.0000"
            )

        with gr.Row():
            with gr.Column():
                comarca = gr.Textbox(label="Comarca", placeholder="Ex: São Paulo")
                vara = gr.Textbox(label="Vara", placeholder="Ex: 1ª")
            with gr.Column():
                cliente = gr.Textbox(label="Nome do Cliente", placeholder="Ex: João da Silva")
                oab = gr.Textbox(label="OAB", placeholder="Ex: SP123456")
        
        fatos = gr.Textbox(
            label="Descrição dos Fatos",
            placeholder="Descreva detalhadamente a situação do cliente...",
            lines=5
        )
        
        testemunhas = gr.Textbox(
            label="Rol de Testemunhas (opcional)",
            placeholder="Nome e qualificação das testemunhas, uma por linha",
            lines=3
        )
        
        output = gr.TextArea(
            label="Peça Gerada",
            lines=30,
            show_copy_button=True
        )

        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",
                "1. Maria Oliveira - Testemunha ocular\n2. José Santos - Empregador do acusado"
            ],
        ], inputs=[tipo, fatos, comarca, vara, cliente, oab, num_processo, testemunhas])

        gr.Button("Gerar Peça").click(
            fn=process,
            inputs=[tipo, fatos, comarca, vara, cliente, oab, num_processo, testemunhas],
            outputs=output
        )

    return demo

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