Update app.py

app.py

@@ -1,42 +1,18 @@
 import gradio as gr
-from datetime import datetime
 from huggingface_hub import InferenceClient
 import os
 
-class DefenseGenerator:
-    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)
-        
-        self.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"
-        }
+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"
+}
 
-    def generate_content(self, prompt: str) -> str:
-        try:
-            response = self.client.chat.completions.create(
-                model="mistralai/Mistral-7B-Instruct-v0.3",
-                messages=[
-                    {"role": "system", "content": "Você é um advogado criminalista experiente. Gere peças sem usar asteriscos ou marcadores. Use numeração romana para seções principais (I, II, III...) e números arábicos para subitens (1, 2, 3...). Mantenha formatação consistente e profissional."},
-                    {"role": "user", "content": prompt}
-                ],
-                max_tokens=2000,
-                temperature=0.7
-            )
-            return response.choices[0].message["content"]
-        except Exception as e:
-            return f"Erro na geração: {str(e)}"
-
-    def gerar_documento(self, tipo: str, fatos: str, dados_basicos: dict) -> str:
-        prompt = f"""Gere uma peça de {self.document_types[tipo]} completa, formatada e sem asteriscos, com base nos seguintes fatos:
+PROMPT_TEMPLATE = """Gere uma peça de {doc_type} completa, formatada e sem asteriscos, com base nos seguintes fatos:
 
 {fatos}
 
@@ -47,150 +23,73 @@ EXCELENTÍSSIMO(A) SENHOR(A) DOUTOR(A) [AUTORIDADE]
 [Título da Peça em Maiúsculas]
 
 I - QUALIFICAÇÃO DAS PARTES
-[Nome e dados completos]
-
 II - FUNDAMENTOS JURÍDICOS
-1. [Primeiro fundamento]
-2. [Segundo fundamento]
-[...]
-
 III - JURISPRUDÊNCIA APLICÁVEL
-1. [Primeira jurisprudência]
-2. [Segunda jurisprudência]
-[...]
-
 IV - DOUTRINA APLICÁVEL
-[Citações doutrinárias relevantes]
-
 V - PRELIMINARES
-1. [Primeira preliminar]
-2. [Segunda preliminar]
-[...]
-
 VI - MÉRITO
-[Desenvolvimento dos argumentos]
-
 VII - PEDIDOS
-1. [Primeiro pedido]
-2. [Segundo pedido]
-[...]
-
-Local, Data.
-
-[Advogado]
-OAB/[Estado] [Número]
 
 Dados para formatação:
-Comarca: {dados_basicos['comarca']}
-Vara: {dados_basicos['vara']}
-Cliente: {dados_basicos['cliente']}
-OAB: {dados_basicos['oab']}"""
+Comarca: {comarca}
+Vara: {vara}
+Cliente: {cliente}
+OAB: {oab}"""
 
-        return self.generate_content(prompt)
+class DefenseGenerator:
+    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 generate_content(self, doc_type: str, fatos: str, dados: dict) -> str:
+        try:
+            prompt = PROMPT_TEMPLATE.format(
+                doc_type=DOCUMENT_TYPES[doc_type],
+                fatos=fatos,
+                **dados
+            )
+            response = self.client.chat.completions.create(
+                model="mistralai/Mistral-7B-Instruct-v0.3",
+                messages=[
+                    {"role": "system", "content": "Você é um advogado criminalista experiente. Gere peças sem usar asteriscos ou marcadores."},
+                    {"role": "user", "content": prompt}
+                ],
+                max_tokens=2000,
+                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):
+        return DefenseGenerator().generate_content(tipo, fatos, 
+            {"comarca": comarca, "vara": vara, "cliente": cliente, "oab": oab})
+
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# Sistema Criminal.AI", elem_classes="text-center text-2xl font-bold mb-4")
+        gr.Markdown("# Sistema Criminal.AI")
         
-        with gr.Group():
-            tipo = gr.Dropdown(
-                choices=list(DefenseGenerator().document_types.keys()),
-                value="resposta_acusacao",
-                label="Tipo de Peça",
-                container=False,
-                scale=1
-            )
+        tipo = gr.Dropdown(choices=list(DOCUMENT_TYPES.keys()), value="resposta_acusacao", label="Tipo de Peça")
         
         with gr.Row():
-            with gr.Column(scale=1):
-                comarca = gr.Textbox(
-                    label="Comarca",
-                    placeholder="Ex: São Paulo",
-                    container=False
-                )
-                vara = gr.Textbox(
-                    label="Vara",
-                    placeholder="Ex: 1ª",
-                    container=False
-                )
-            
-            with gr.Column(scale=1):
-                cliente = gr.Textbox(
-                    label="Nome do Cliente",
-                    placeholder="Ex: João da Silva",
-                    container=False
-                )
-                oab = gr.Textbox(
-                    label="OAB",
-                    placeholder="Ex: SP123456",
-                    container=False
-                )
-        
-        fatos = gr.Textbox(
-            label="Descreva os fatos",
-            placeholder="Descreva a situação do cliente e o que precisa ser argumentado...",
-            lines=5,
-            container=False
-        )
+            comarca = gr.Textbox(label="Comarca", placeholder="Ex: São Paulo")
+            vara = gr.Textbox(label="Vara", placeholder="Ex: 1ª")
+            cliente = gr.Textbox(label="Nome do Cliente", placeholder="Ex: João da Silva")
+            oab = gr.Textbox(label="OAB", placeholder="Ex: SP123456")
         
-        with gr.Row():
-            gerar_btn = gr.Button("Gerar Peça", variant="primary", scale=2)
-            limpar_btn = gr.Button("Limpar", scale=1)
-            
-        output = gr.TextArea(
-            label="Documento Gerado",
-            lines=20,
-            show_copy_button=True,
-            container=False
-        )
-
-        def process(tipo, fatos, comarca, vara, cliente, oab):
-            gerador = DefenseGenerator()
-            dados = {
-                "comarca": comarca,
-                "vara": vara,
-                "cliente": cliente,
-                "oab": oab
-            }
-            return gerador.gerar_documento(tipo, fatos, dados)
-            
-        def clear():
-            return {
-                tipo: "resposta_acusacao",
-                comarca: "",
-                vara: "",
-                cliente: "",
-                oab: "",
-                fatos: "",
-                output: ""
-            }
-            
-        gerar_btn.click(
-            fn=process,
-            inputs=[tipo, fatos, comarca, vara, cliente, oab],
-            outputs=output,
-            api_name="gerar_peca"
-        )
+        fatos = gr.Textbox(label="Descreva os fatos", lines=5)
+        output = gr.TextArea(label="Documento Gerado", lines=20, show_copy_button=True)
         
-        limpar_btn.click(
-            fn=clear,
-            outputs=[tipo, comarca, vara, cliente, oab, fatos, output],
-            api_name="limpar"
-        )
+        gr.Examples([
+            ["resposta_acusacao", "Cliente preso em flagrante por furto de celular iPhone 13.", "São Paulo", "1ª", "João Silva", "SP123456"],
+            ["habeas_corpus", "Cliente preso preventivamente há 90 dias por tráfico.", "São Paulo", "2ª", "Pedro Santos", "SP654321"]
+        ], inputs=[tipo, fatos, comarca, vara, cliente, oab])
 
-        exemplos = [
-            ["resposta_acusacao", "Cliente preso em flagrante por furto de celular iPhone 13 da vítima Maria em 10/01/2024. É primário, tem residência fixa e trabalho. A vítima não reconheceu com certeza.", "São Paulo", "1ª", "João Silva", "SP123456"],
-            ["habeas_corpus", "Cliente preso preventivamente há 90 dias por tráfico. Apreendidas 10g de maconha, é primário, tem residência e trabalho fixos. Não foram apreendidos outros materiais típicos de tráfico.", "São Paulo", "2ª", "Pedro Santos", "SP654321"]
-        ]
-        
-        gr.Examples(
-            examples=exemplos,
-            inputs=[tipo, fatos, comarca, vara, cliente, oab],
-            label="Exemplos de Uso"
-        )
+        gr.Button("Gerar Peça").click(fn=process, inputs=[tipo, fatos, comarca, vara, cliente, oab], outputs=output)
 
     return demo
 
 if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch()
+    create_interface().launch()