Update app.py

app.py

@@ -12,74 +12,121 @@ 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):
-    """
-    Gera uma peça jurídica criminal em português do Brasil com base no tipo e nas informações fornecidas.
-
-    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 com formatação adequada.
-    """
-    # Verifica se o tipo de peça foi selecionado
     if not tipo_peca:
         return "Por favor, selecione o tipo de peça jurídica."
 
-    # Mensagens para o modelo
-    mensagens_iniciais = [
+    # 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 especialista em direito criminal no Brasil e deve gerar peças jurídicas detalhadas. "
-                "Todas as respostas devem estar em português do Brasil, considerando o ordenamento jurídico brasileiro."
+                "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"Quero que você elabore um(a) {tipo_peca} com base nas seguintes informações do caso: {informacoes_caso}"
+            "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'])}"
         }
     ]
 
-    # Primeira requisição: Geração inicial da peça
-    resposta_inicial = client.chat.completions.create(
+    resposta_estrutura = client.chat.completions.create(
         model="mistralai/Mistral-7B-Instruct-v0.3",
-        messages=mensagens_iniciais,
-        max_tokens=1000
+        messages=mensagens_estrutura,
+        max_tokens=1500,
+        temperature=0.7
     )
-    
-    peca_inicial = resposta_inicial.choices[0].message["content"]
 
-    # Segunda requisição: Expansão da peça jurídica com mais detalhes e argumentos
-    mensagens_expansao = [
+    estrutura_base = resposta_estrutura.choices[0].message["content"]
+
+    # Segunda chamada: Aprofundamento legal
+    mensagens_legal = [
         {
             "role": "system",
             "content": (
-                "Você é um especialista em direito criminal no Brasil e deve expandir a peça jurídica gerada, "
-                "adicionando mais detalhes, justificativas e argumentos legais. "
-                "Todas as respostas devem estar em português do Brasil, considerando o ordenamento jurídico brasileiro."
+                "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"Agora, com base na peça inicial a seguir, adicione mais detalhes, justificativas e argumentos legais:\n{peca_inicial}"
+            "content": f"Aprimore esta peça com fundamentação legal:\n{estrutura_base}"
         }
     ]
 
-    # Segunda requisição: Expansão e aprofundamento da peça jurídica
-    resposta_expansao = client.chat.completions.create(
+    resposta_legal = client.chat.completions.create(
         model="mistralai/Mistral-7B-Instruct-v0.3",
-        messages=mensagens_expansao,
-        max_tokens=1000
+        messages=mensagens_legal,
+        max_tokens=1500,
+        temperature=0.7
     )
 
-    peca_expandida = resposta_expansao.choices[0].message["content"]
+    # 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']}"
+            )
+        }
+    ]
 
-    # Combina a peça inicial com a expansão, garantindo um documento completo e coerente
-    peca_completa = f"**Peça Jurídica: {tipo_peca.capitalize()}**\n\n{peca_inicial}\n\n{peca_expandida}"
+    resposta_final = client.chat.completions.create(
+        model="mistralai/Mistral-7B-Instruct-v0.3",
+        messages=mensagens_revisao,
+        max_tokens=1500,
+        temperature=0.5
+    )
 
-    # Retorna a peça completa
-    return peca_completa
+    return resposta_final.choices[0].message["content"]