Update app.py

app.py

@@ -1,218 +1,84 @@
 import os
 import gradio as gr
-from huggingface_hub import HfApi, InferenceApi
-from datetime import datetime
-import logging
-from typing import Dict, List
+from huggingface_hub import InferenceClient
 
-# Configuração do logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# 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.")
 
-class DocumentGenerator:
-    """Gerencia a geração de documentos usando BERT-PT"""
-    
-    def __init__(self):
-        # Usando modelo BERT português
-        self.client = InferenceApi(
-            repo_id="neuralmind/bert-base-portuguese-legal-cased",
-            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"""[DOCUMENTO LEGAL]
-TIPO: {doc_type}
-
-[QUALIFICAÇÃO]
-CLIENTE: {context.get('client_name')}
-PROCESSO: {context.get('process_number')}
-TRIBUNAL: {context.get('court')}
-COMARCA: {context.get('jurisdiction')}
-
-[FATOS]
-{context.get('facts')}
+# Criação do cliente com a API Key
+client = InferenceClient(api_key=api_key)
 
-[FUNDAMENTOS]
-{context.get('legal_basis')}
+# Função para gerar a peça jurídica em português
+def gerar_peca_juridica_criminal(informacoes_caso, tipo_peca):
+    """
+    Gera uma peça jurídica criminal em português com base no tipo e nas informações fornecidas.
 
-[INSTRUÇÕES]
-- Use linguagem jurídica formal
-- Siga formato brasileiro
-- Inclua todas as seções necessárias
-- Mantenha clareza e precisão
-- Use formatação adequada
+    Parâmetros:
+        informacoes_caso (str): Descrição detalhada do caso.
+        tipo_peca (str): Tipo de peça jurídica a ser gerada (ex: habeas corpus, petição inicial).
+    
+    Retorna:
+        str: Peça jurídica gerada.
+    """
+    mensagens = [
+        {
+            "role": "system",
+            "content": (
+                "Você é um especialista em direito criminal no Brasil e deve gerar peças jurídicas detalhadas. "
+                "Todas as respostas devem estar em português, considerando o ordenamento jurídico brasileiro."
+            )
+        },
+        {
+            "role": "user",
+            "content": f"Quero que você elabore um(a) {tipo_peca} com base nas seguintes informações do caso: {informacoes_caso}"
+        }
+    ]
 
-Gere o documento completo em português:"""
+    # Chamada ao modelo para gerar a resposta
+    resposta = client.chat.completions.create(
+        model="mistralai/Mistral-7B-Instruct-v0.3",
+        messages=mensagens,
+        max_tokens=1000
+    )
 
-            # Faz a chamada ao modelo
-            response = self.client.text_generation(
-                prompt,
-                max_new_tokens=1024,
-                do_sample=True,
-                temperature=0.2,
-                top_p=0.95,
-                num_return_sequences=1,
-                repetition_penalty=1.2
-            )
+    return resposta.choices[0].message["content"]
 
-            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
-        text = text.split("Gere o documento completo em português:")[-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
-        current_date = datetime.now().strftime('%d de %B de %Y').replace('May', 'Maio').replace('April', 'Abril')
-        formatted_text = f"{formatted_text}\n\n{context.get('jurisdiction')}, {current_date}"
-        
-        return formatted_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"""
+# Configuração do app Gradio
+def gerar_interface_gradio():
+    """
+    Configura a interface do Gradio para gerar peças jurídicas.
+    """
+    with gr.Blocks() as interface:
+        gr.Markdown("## Gerador de Peças Jurídicas Criminais")
+        gr.Markdown(
+            "Insira as informações detalhadas do caso e escolha o tipo de peça jurídica que deseja gerar."
+        )
         
-        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
-                ]
+        with gr.Row():
+            entrada_caso = gr.Textbox(
+                label="Informações do Caso",
+                placeholder="Descreva o caso detalhadamente aqui...",
+                lines=5
             )
-            
-            # Eventos
-            generate_btn.click(
-                fn=self._generate_document,
-                inputs=[
-                    doc_type, client_name, process_number,
-                    court, jurisdiction, facts, legal_basis
-                ],
-                outputs=output
+            entrada_tipo = gr.Textbox(
+                label="Tipo de Peça Jurídica",
+                placeholder="Exemplo: habeas corpus, petição inicial..."
             )
-    
-    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)}"
+        botao_gerar = gr.Button("Gerar Peça Jurídica")
+        saida_texto = gr.Textbox(label="Peça Jurídica Gerada", lines=15)
+
+        botao_gerar.click(
+            gerar_peca_juridica_criminal, 
+            inputs=[entrada_caso, entrada_tipo], 
+            outputs=saida_texto
+        )
     
-    def launch(self):
-        """Inicia a interface web"""
-        self.app.launch(share=True)
+    return interface
 
+# Executa o app
 if __name__ == "__main__":
-    interface = WebInterface()
-    interface.launch()
+    app = gerar_interface_gradio()
+    app.launch()