Update app.py

app.py

@@ -11,91 +11,84 @@ if not 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):
+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."
 
-    template = {
-        "estrutura": [
-            "EXCELENTÍSSIMO(A) SENHOR(A) JUIZ(A)",
-            "FATOS",
-            "DIREITO",
-            "PEDIDOS"
-        ],
-        "prompt": (
-            "Gere HC técnico e conciso:\n"
-            "1. Qualificação e fatos\n"
-            "2. Fundamento: constrangimento ilegal\n"
-            "3. Pedido liminar: risco e fundamento\n"
-            "4. Pedido final: soltura"
-        )
-    }
+    if tipo_peca.lower() != "habeas corpus":
+        return "Tipo de peça não suportado no momento. Disponível apenas 'habeas corpus'."
 
+    # Template com limite de tokens
     mensagens_estrutura = [
         {
             "role": "system",
-            "content": (
-                "Como advogado criminalista brasileiro, redija um HC conciso e técnico em português formal.\n"
-                "1. Use numeração de parágrafos\n"
-                "2. Evite repetições\n"
-                "3. Cite corretamente apenas legislação existente\n"
-                "4. Formule pedidos com precisão técnica"
-            )
+            "content": "Redija HC técnico em português jurídico. Máximo 300 tokens."
         },
         {
             "role": "user",
-            "content": f"{template['prompt']}\nCaso: {informacoes_caso}"
+            "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}"
+            )
         }
     ]
 
     resposta_estrutura = client.chat.completions.create(
         model="mistralai/Mistral-7B-Instruct-v0.3",
         messages=mensagens_estrutura,
-        max_tokens=1000,
+        max_tokens=300,
         temperature=0.7
     )
 
     estrutura_base = resposta_estrutura.choices[0].message["content"]
 
+    # Elaboração com limite
     mensagens_elaboracao = [
         {
             "role": "system",
-            "content": "Como advogado criminalista brasileiro, fundamente cada argumento com legislação e jurisprudência pertinente."
+            "content": "Aprimore com fundamentação técnica. Máximo 300 tokens."
         },
         {
             "role": "user",
-            "content": f"Aprimore esta peça mantendo os argumentos essenciais:\n{estrutura_base}"
+            "content": f"Revise mantendo concisão:\n{estrutura_base}"
         }
     ]
 
     resposta_elaboracao = client.chat.completions.create(
         model="mistralai/Mistral-7B-Instruct-v0.3",
         messages=mensagens_elaboracao,
-        max_tokens=1000,
+        max_tokens=300,
         temperature=0.7
     )
 
+    # Finalização com limite
     mensagens_final = [
         {
             "role": "system",
-            "content": (
-                "Faça revisão técnica final garantindo:\n"
-                "1. Coesão e clareza\n"
-                "2. Embasamento legal preciso\n"
-                "3. Pedidos bem fundamentados\n"
-                "4. Linguagem jurídica correta"
-            )
+            "content": "Revise tecnicamente. Máximo 400 tokens."
         },
         {
             "role": "user",
-            "content": f"Elabore versão final técnica:\n{estrutura_base}\n\n{resposta_elaboracao.choices[0].message['content']}"
+            "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']}"
+            )
         }
     ]
 
     resposta_final = client.chat.completions.create(
         model="mistralai/Mistral-7B-Instruct-v0.3",
         messages=mensagens_final,
-        max_tokens=1000,
+        max_tokens=400,
         temperature=0.5
     )