app.py

import os
import gradio as gr
from huggingface_hub import InferenceClient

# Obtendo a API Key da variável de ambiente
api_key = os.getenv("HF_API_KEY")
if not api_key:
    raise EnvironmentError("A variável de ambiente 'HF_API_KEY' não está configurada.")

# Criação do cliente com a API Key
client = InferenceClient(api_key=api_key)

# Função para gerar a peça jurídica em português do Brasil
def gerar_peca_juridica_criminal(informacoes_caso, tipo_peca):
    if not tipo_peca:
        return "Por favor, selecione o tipo de peça jurídica."

    # Template específico para cada tipo de peça
    templates = {
        "habeas corpus": {
            "estrutura": [
                "EXCELENTÍSSIMO(A) SENHOR(A) DOUTOR(A) JUIZ(A) DE DIREITO",
                "FUNDAMENTOS JURÍDICOS",
                "PEDIDOS"
            ],
            "prompt": "Gere um Habeas Corpus detalhado considerando:\n"
                     "1. Fundamentos constitucionais (Art. 5º, LXVIII, CF)\n"
                     "2. Requisitos específicos do HC\n"
                     "3. Jurisprudência relevante do STF/STJ\n"
                     "4. Pedidos liminares e principais\n"
        },
        "petição inicial": {
            "estrutura": [
                "QUALIFICAÇÃO DAS PARTES",
                "DOS FATOS",
                "DO DIREITO",
                "DOS PEDIDOS"
            ],
            "prompt": "Elabore uma petição inicial criminal contemplando:\n"
                     "1. Elementos essenciais (Art. 41, CPP)\n"
                     "2. Fundamentação legal específica\n"
                     "3. Provas e documentos relevantes\n"
                     "4. Pedidos principais e subsidiários\n"
        }
    }

    template = templates.get(tipo_peca.lower(), {
        "estrutura": ["INTRODUÇÃO", "DESENVOLVIMENTO", "CONCLUSÃO"],
        "prompt": f"Elabore uma {tipo_peca} criminal completa e fundamentada."
    })

    # Primeira chamada: Estruturação básica
    mensagens_estrutura = [
        {
            "role": "system",
            "content": (
                "Você é um advogado criminal especialista brasileiro. "
                "Formate a peça seguindo rigorosamente a estrutura jurídica brasileira. "
                "Use linguagem técnica apropriada e citações legais precisas."
            )
        },
        {
            "role": "user",
            "content": f"{template['prompt']}\n\nCaso: {informacoes_caso}\n\n"
                      f"Estruture a peça com as seguintes seções:\n"
                      f"{chr(10).join(template['estrutura'])}"
        }
    ]

    resposta_estrutura = client.chat.completions.create(
        model="mistralai/Mistral-7B-Instruct-v0.3",
        messages=mensagens_estrutura,
        max_tokens=1500,
        temperature=0.7
    )

    estrutura_base = resposta_estrutura.choices[0].message["content"]

    # Segunda chamada: Aprofundamento legal
    mensagens_legal = [
        {
            "role": "system",
            "content": (
                "Você é um especialista em direito penal brasileiro. "
                "Analise a peça e adicione fundamentação legal robusta, "
                "incluindo doutrina, jurisprudência e legislação específica."
            )
        },
        {
            "role": "user",
            "content": f"Aprimore esta peça com fundamentação legal:\n{estrutura_base}"
        }
    ]

    resposta_legal = client.chat.completions.create(
        model="mistralai/Mistral-7B-Instruct-v0.3",
        messages=mensagens_legal,
        max_tokens=1500,
        temperature=0.7
    )

    # Terceira chamada: Revisão e formatação final
    mensagens_revisao = [
        {
            "role": "system",
            "content": (
                "Você é um revisor jurídico expert. "
                "Faça a revisão final garantindo clareza, coerência "
                "e formatação adequada da peça."
            )
        },
        {
            "role": "user",
            "content": (
                f"Revise e formate esta peça jurídica:\n"
                f"{estrutura_base}\n\n"
                f"Complementos legais:\n{resposta_legal.choices[0].message['content']}"
            )
        }
    ]

    resposta_final = client.chat.completions.create(
        model="mistralai/Mistral-7B-Instruct-v0.3",
        messages=mensagens_revisao,
        max_tokens=1500,
        temperature=0.5
    )

    return resposta_final.choices[0].message["content"]



# Função para redefinir os campos
def redefinir_campos():
    return "", "", ""

# Função para exportar a peça jurídica gerada
def exportar_peca(peca_texto):
    """
    Exporta a peça jurídica gerada para um arquivo .txt.

    Parâmetros:
        peca_texto (str): Texto da peça jurídica gerada.
    
    Retorna:
        str: Caminho do arquivo exportado.
    """
    arquivo_nome = "peca_juridica_gerada.txt"
    with open(arquivo_nome, "w", encoding="utf-8") as arquivo:
        arquivo.write(peca_texto)
    return f"Arquivo salvo como: {arquivo_nome}"

# Configuração do app Gradio com funcionalidades adicionais
def gerar_interface_gradio():
    """
    Configura a interface do Gradio para gerar peças jurídicas com exemplos pré-definidos e funcionalidades adicionais.
    """
    with gr.Blocks() as interface:
        gr.Markdown("# Gerador de Peças Jurídicas Criminais")
        gr.Markdown(
            "### Preencha os campos abaixo ou escolha um exemplo, e selecione o tipo de peça jurídica que deseja gerar."
        )
        
        with gr.Row():
            entrada_caso = gr.Textbox(
                label="Informações do Caso",
                placeholder="Descreva o caso detalhadamente ou selecione um exemplo abaixo.",
                lines=5
            )
            
            exemplo_caso = gr.Dropdown(
                label="Exemplos de Casos",
                choices=[
                    "O cliente foi acusado injustamente de furto qualificado, mas há provas de que estava em outro local no momento do crime.",
                    "O cliente foi preso em flagrante por porte de entorpecentes, mas os policiais não apresentaram o mandado de busca.",
                    "O cliente foi acusado de agressão, mas existem testemunhas que comprovam sua ausência no local do ocorrido."
                ],
                value=None,
                interactive=True
            )
            
            entrada_tipo = gr.Dropdown(
                label="Tipo de Peça Jurídica",
                choices=["habeas corpus", "petição inicial", "defesa prévia", "apelação criminal"],
                value=None  # Define o valor inicial como nenhum
            )
        
        def atualizar_caso(caso_exemplo):
            """
            Atualiza o campo de texto com o exemplo escolhido pelo usuário.
            """
            return caso_exemplo

        exemplo_caso.change(
            atualizar_caso, 
            inputs=exemplo_caso, 
            outputs=entrada_caso
        )
        
        with gr.Row():
            botao_gerar = gr.Button("Gerar Peça Jurídica")
            botao_redefinir = gr.Button("Redefinir Campos")
        
        saida_texto = gr.Textbox(
            label="Peça Jurídica Gerada",
            lines=15,
            placeholder="A peça jurídica gerada aparecerá aqui."
        )

        botao_gerar.click(
            gerar_peca_juridica_criminal, 
            inputs=[entrada_caso, entrada_tipo], 
            outputs=saida_texto
        )

        botao_redefinir.click(
            redefinir_campos, 
            inputs=None, 
            outputs=[entrada_caso, entrada_tipo, saida_texto]
        )
        
        with gr.Row():
            botao_exportar = gr.Button("Exportar Peça Jurídica para .txt")
            saida_arquivo = gr.Textbox(
                label="Arquivo Exportado",
                interactive=False,
                placeholder="O caminho do arquivo exportado será exibido aqui."
            )
        
        botao_exportar.click(
            exportar_peca,
            inputs=saida_texto,
            outputs=saida_arquivo
        )
    
    return interface

# Executa o app
if __name__ == "__main__":
    app = gerar_interface_gradio()
    app.launch()