Update app.py

app.py

@@ -15,86 +15,115 @@ def gerar_peca_juridica_criminal(informacoes_caso: str, tipo_peca: str) -> str:
     if not tipo_peca or not informacoes_caso:
         return "Por favor, forneça todas as informações necessárias."
 
-    if tipo_peca.lower() != "habeas corpus":
-        return "Tipo de peça não suportado no momento. Disponível apenas 'habeas corpus'."
+    templates = {
+        "habeas corpus": {
+            "estrutura": ["EXORDIAL", "FATOS", "DIREITO", "PEDIDOS"],
+            "prompt": (
+                "HC conciso:\n"
+                "1. Fatos e constrangimento ilegal\n"
+                "2. CF art 5º LXVIII e CPP 647/648\n"
+                "3. Precedente STF/STJ\n"
+                "4. Liminar + principal"
+            ),
+            "max_tokens": 300
+        },
+        "petição inicial": {
+            "estrutura": ["EXORDIAL", "QUALIFICAÇÃO", "FATOS", "DIREITO", "PEDIDOS"],
+            "prompt": (
+                "Denúncia criminal:\n"
+                "1. Qualificação completa\n"
+                "2. Fatos e tipificação\n"
+                "3. Arts. 41/395 CPP\n"
+                "4. Provas e diligências"
+            ),
+            "max_tokens": 400
+        },
+        "defesa prévia": {
+            "estrutura": ["QUALIFICAÇÃO", "PRELIMINARES", "MÉRITO", "PROVAS", "PEDIDOS"],
+            "prompt": (
+                "Defesa art. 396-A CPP:\n"
+                "1. Preliminares\n"
+                "2. Tese defensiva\n"
+                "3. Provas a produzir\n"
+                "4. Testemunhas"
+            ),
+            "max_tokens": 350
+        },
+        "apelação criminal": {
+            "estrutura": ["TEMPESTIVIDADE", "FATOS", "PRELIMINARES", "MÉRITO", "PEDIDOS"],
+            "prompt": (
+                "Apelação criminal:\n"
+                "1. Art. 593 CPP\n" 
+                "2. Preliminares\n"
+                "3. Error in judicando\n"
+                "4. Reforma/nulidade"
+            ),
+            "max_tokens": 450
+        }
+    }
+
+    template = templates.get(tipo_peca.lower())
+    if not template:
+        return "Tipo de peça não suportado."
 
-    # Template com limite de tokens
     mensagens_estrutura = [
         {
             "role": "system",
-            "content": "Redija HC técnico em português jurídico. Máximo 300 tokens."
+            "content": f"Redija {tipo_peca} técnica em português jurídico. Máximo {template['max_tokens']} tokens."
         },
         {
             "role": "user",
-            "content": (
-                f"Elabore HC conciso com:\n"
-                f"1. Qualificação e fatos essenciais\n"
-                f"2. Art 5º LXVIII CF e art 647/648 CPP\n"
-                f"3. Um precedente STF/STJ\n"
-                f"4. Pedido liminar e principal\n\n"
-                f"Caso: {informacoes_caso}"
-            )
+            "content": f"{template['prompt']}\nCaso: {informacoes_caso}"
         }
     ]
 
     resposta_estrutura = client.chat.completions.create(
         model="mistralai/Mistral-7B-Instruct-v0.3",
         messages=mensagens_estrutura,
-        max_tokens=300,
+        max_tokens=template['max_tokens'],
         temperature=0.7
     )
 
     estrutura_base = resposta_estrutura.choices[0].message["content"]
 
-    # Elaboração com limite
     mensagens_elaboracao = [
         {
             "role": "system",
-            "content": "Aprimore com fundamentação técnica. Máximo 300 tokens."
+            "content": f"Aprimore {tipo_peca} com fundamentação precisa. Máximo {template['max_tokens']} tokens."
         },
         {
             "role": "user",
-            "content": f"Revise mantendo concisão:\n{estrutura_base}"
+            "content": f"Revise:\n{estrutura_base}"
         }
     ]
 
     resposta_elaboracao = client.chat.completions.create(
         model="mistralai/Mistral-7B-Instruct-v0.3",
         messages=mensagens_elaboracao,
-        max_tokens=300,
+        max_tokens=template['max_tokens'],
         temperature=0.7
     )
 
-    # Finalização com limite
     mensagens_final = [
         {
             "role": "system",
-            "content": "Revise tecnicamente. Máximo 400 tokens."
+            "content": f"Finalize {tipo_peca} em formato técnico. Estrutura: {' > '.join(template['estrutura'])}"
         },
         {
             "role": "user",
-            "content": (
-                f"Finalize respeitando estrutura:\n"
-                "EXCELENTÍSSIMO JUIZ\n"
-                "FATOS\n"
-                "DIREITO\n"
-                "PEDIDOS\n\n"
-                f"{estrutura_base}\n\n"
-                f"{resposta_elaboracao.choices[0].message['content']}"
-            )
+            "content": f"Versão final:\n{estrutura_base}\n\n{resposta_elaboracao.choices[0].message['content']}"
         }
     ]
 
     resposta_final = client.chat.completions.create(
         model="mistralai/Mistral-7B-Instruct-v0.3",
         messages=mensagens_final,
-        max_tokens=400,
+        max_tokens=max(template['max_tokens'] + 100, 500),
         temperature=0.5
     )
 
     return resposta_final.choices[0].message["content"]
     
-    
 # Função para redefinir os campos
 def redefinir_campos():
     return "", "", ""