app.py

import gradio as gr
from gradio_client import Client
import json
import sqlite3
from datetime import datetime
import os
from typing import Dict, List, Optional
import asyncio
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class DatabaseManager:
    def __init__(self, db_path: str = "legal_docs.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS documents (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    title TEXT,
                    doc_type TEXT,
                    content TEXT,
                    created_at TIMESTAMP,
                    client_name TEXT,
                    process_number TEXT,
                    template_used TEXT
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS templates (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT,
                    doc_type TEXT,
                    content TEXT,
                    created_at TIMESTAMP
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS jurisprudence (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    topic TEXT,
                    content TEXT,
                    source TEXT,
                    relevance INTEGER
                )
            """)

class PromptOptimizer:
    def __init__(self):
        self.base_prompts = self._load_base_prompts()
        
    def _load_base_prompts(self) -> Dict[str, str]:
        return {
            "habeas_corpus": """
            Atue como um advogado criminalista experiente e gere um Habeas Corpus detalhado e tecnicamente preciso:
            
            CONTEXTO DO CASO:
            {caso_detalhes}
            
            REQUISITOS ESPECÍFICOS:
            1. Use linguagem jurídica formal
            2. Cite jurisprudência relevante dos tribunais superiores
            3. Estruture o documento com todas as partes necessárias
            4. Inclua argumentos constitucionais pertinentes
            5. Faça referências específicas ao CPP
            
            ESTRUTURA REQUERIDA:
            - Cabeçalho completo com qualificação
            - Fatos detalhados
            - Fundamentos jurídicos
            - Jurisprudência aplicável
            - Pedidos liminar e final
            """,
            
            "denuncia_crime": """
            Atue como um advogado criminalista experiente e gere uma denúncia criminal detalhada:
            
            DETALHES DO CASO:
            {caso_detalhes}
            
            REQUISITOS:
            1. Descreva os fatos com precisão temporal
            2. Tipifique corretamente a conduta
            3. Indique as circunstâncias do crime
            4. Qualifique todos os envolvidos
            5. Indique as provas disponíveis
            
            ESTRUTURA:
            - Qualificação completa
            - Narrativa dos fatos
            - Tipificação penal
            - Pedidos
            - Provas a produzir
            """,
            
            "recurso_especial": """
            Atue como um advogado criminalista e gere um Recurso Especial Criminal:
            
            CASO E DECISÃO RECORRIDA:
            {caso_detalhes}
            
            REQUISITOS:
            1. Demonstre o prequestionamento
            2. Indique violação à lei federal
            3. Cite precedentes do STJ
            4. Demonstre divergência jurisprudencial
            
            ESTRUTURA:
            - Cabimento e tempestividade
            - Prequestionamento
            - Mérito recursal
            - Pedidos
            """
        }
    
    def optimize_prompt(self, doc_type: str, context: Dict[str, str]) -> str:
        base_prompt = self.base_prompts.get(doc_type, "")
        optimized = base_prompt.format(caso_detalhes=json.dumps(context, indent=2))
        return optimized

class LegalDocumentGenerator:
    def __init__(self):
        self.client = Client("ysharma/Chat_with_Meta_llama3_8b")
        self.db = DatabaseManager()
        self.prompt_optimizer = PromptOptimizer()
        
    async def generate_document(self, doc_type: str, context: Dict[str, str], 
                              template_id: Optional[int] = None) -> str:
        try:
            # Otimiza o prompt baseado no tipo de documento
            optimized_prompt = self.prompt_optimizer.optimize_prompt(doc_type, context)
            
            # Gera o documento usando a API do Llama
            result = await self._generate_with_llama(optimized_prompt)
            
            # Salva o documento gerado
            self._save_document(doc_type, result, context)
            
            return result
            
        except Exception as e:
            logger.error(f"Erro na geração do documento: {str(e)}")
            raise

    async def _generate_with_llama(self, prompt: str) -> str:
        return self.client.predict(
            message=prompt,
            request=0.85,  # Temperatura reduzida para maior consistência
            param_3=3072,  # Aumentado para documentos mais longos
            api_name="/chat"
        )
    
    def _save_document(self, doc_type: str, content: str, context: Dict[str, str]):
        with sqlite3.connect(self.db.db_path) as conn:
            conn.execute("""
                INSERT INTO documents 
                (title, doc_type, content, created_at, client_name, process_number)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (
                f"{doc_type}_{context.get('client_name')}",
                doc_type,
                content,
                datetime.now(),
                context.get('client_name'),
                context.get('process_number')
            ))

class AdvancedLegalInterface:
    def __init__(self):
        self.generator = LegalDocumentGenerator()
        self.create_interface()
    
    def create_interface(self):
        with gr.Blocks(theme=gr.themes.Glass()) as self.app:
            gr.Markdown("""
            # Sistema Avançado de Geração de Peças Criminais
            ### Powered by LLM Technology
            """)
            
            with gr.Tabs():
                with gr.Tab("Gerar Documento"):
                    with gr.Row():
                        with gr.Column():
                            doc_type = gr.Dropdown(
                                choices=[
                                    "Habeas Corpus",
                                    "Denúncia Criminal",
                                    "Recurso Especial",
                                    "Alegações Finais",
                                    "Resposta à Acusação",
                                    "Revisão Criminal",
                                    "Apelação Criminal"
                                ],
                                label="Tipo de Documento"
                            )
                            
                            client_info = gr.Group()
                            with client_info:
                                client_name = gr.Textbox(label="Nome do Cliente")
                                process_number = gr.Textbox(label="Número do Processo")
                                court = gr.Textbox(label="Tribunal")
                                jurisdiction = gr.Textbox(label="Comarca")
                            
                            case_details = gr.Group()
                            with case_details:
                                facts = gr.Textbox(
                                    label="Fatos",
                                    lines=5,
                                    placeholder="Descreva os fatos relevantes..."
                                )
                                legal_basis = gr.Textbox(
                                    label="Fundamentos Jurídicos",
                                    lines=3,
                                    placeholder="Indique os principais fundamentos legais..."
                                )
                            
                            advanced_options = gr.Group()
                            with advanced_options:
                                use_template = gr.Checkbox(label="Usar Template Personalizado")
                                template_select = gr.Dropdown(
                                    choices=self._get_templates(),
                                    label="Template",
                                    visible=False
                                )
                                include_jurisprudence = gr.Checkbox(
                                    label="Incluir Jurisprudência Relevante"
                                )
                            
                            generate_btn = gr.Button("Gerar Documento", variant="primary")
                            
                        with gr.Column():
                            output = gr.Textbox(
                                label="Documento Gerado",
                                lines=30,
                                show_copy_button=True
                            )
                            
                            with gr.Row():
                                save_btn = gr.Button("Salvar Documento")
                                export_btn = gr.Button("Exportar como PDF")
                
                with gr.Tab("Gerenciar Templates"):
                    self._create_template_manager()
                
                with gr.Tab("Histórico"):
                    self._create_history_viewer()
            
            # Eventos
            generate_btn.click(
                fn=self._generate_document,
                inputs=[
                    doc_type, client_name, process_number, court,
                    jurisdiction, facts, legal_basis, use_template,
                    template_select, include_jurisprudence
                ],
                outputs=output
            )
            
            use_template.change(
                fn=lambda x: gr.update(visible=x),
                inputs=[use_template],
                outputs=[template_select]
            )
    
    def _get_templates(self) -> List[str]:
        with sqlite3.connect(self.generator.db.db_path) as conn:
            templates = conn.execute("SELECT name FROM templates").fetchall()
            return [t[0] for t in templates]
    
    def _create_template_manager(self):
        with gr.Group():
            gr.Markdown("### Gerenciar Templates")
            template_name = gr.Textbox(label="Nome do Template")
            template_content = gr.Textbox(label="Conteúdo", lines=10)
            save_template_btn = gr.Button("Salvar Template")
    
    def _create_history_viewer(self):
        with gr.Group():
            gr.Markdown("### Histórico de Documentos")
            history_table = gr.Dataframe(
                headers=["Data", "Tipo", "Cliente", "Processo"],
                row_count=10
            )
    
    async def _generate_document(self, *args) -> str:
        try:
            context = {
                "client_name": args[1],
                "process_number": args[2],
                "court": args[3],
                "jurisdiction": args[4],
                "facts": args[5],
                "legal_basis": args[6],
                "use_template": args[7],
                "template": args[8] if args[7] else None,
                "include_jurisprudence": args[9]
            }
            
            result = await self.generator.generate_document(
                doc_type=args[0].lower().replace(" ", "_"),
                context=context
            )
            
            return result
        
        except Exception as e:
            logger.error(f"Erro na geração: {str(e)}")
            return f"Erro na geração do documento: {str(e)}"
    
    def launch(self):
        self.app.launch(share=True)

if __name__ == "__main__":
    interface = AdvancedLegalInterface()
    interface.launch()