app.py

"""
🧠 Agent Pruning Lab
Demonstração visual de Context Pruning com Provence e Semantic Highlighting
para sistemas de Agentes LLM - Hugging Face Space
"""

import streamlit as st
import json
import random
from components.visualizer import SemanticHighlighter
from components.provence_wrapper import ProvencePruner
from components.agent_simulator import AgentSimulator

# Configuração da página
st.set_page_config(
    page_title="🧠 Agent Pruning Lab",
    page_icon="✂️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# CSS Customizado
st.markdown("""
<style>
    .main-header {
        font-size: 3rem;
        font-weight: 800;
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        margin-bottom: 0.5rem;
    }
    .sub-header {
        font-size: 1.2rem;
        color: #666;
        margin-bottom: 2rem;
    }
    .metric-card {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        padding: 1.5rem;
        border-radius: 15px;
        color: white;
        text-align: center;
        box-shadow: 0 10px 30px rgba(102, 126, 234, 0.3);
    }
    .metric-value {
        font-size: 2.5rem;
        font-weight: 700;
    }
    .metric-label {
        font-size: 0.9rem;
        opacity: 0.9;
    }
    .pruned-text {
        background-color: #ff6b6b20;
        text-decoration: line-through;
        color: #999;
        padding: 2px 4px;
        border-radius: 3px;
    }
    .kept-text {
        background-color: #51cf6620;
        color: #2b8a3e;
        padding: 2px 4px;
        border-radius: 3px;
        font-weight: 500;
    }
    .highlight-low { background-color: #fff3bf; }
    .highlight-med { background-color: #ffd8a8; }
    .highlight-high { background-color: #ff8787; }
    .agent-thought {
        background: #f8f9fa;
        border-left: 4px solid #667eea;
        padding: 1rem;
        margin: 0.5rem 0;
        border-radius: 0 8px 8px 0;
        font-family: 'Courier New', monospace;
    }
    .stTabs [data-baseweb="tab-list"] {
        gap: 8px;
    }
    .stTabs [data-baseweb="tab"] {
        background-color: #f0f2f6;
        border-radius: 8px 8px 0 0;
        padding: 10px 20px;
    }
</style>
""", unsafe_allow_html=True)

# Inicialização dos componentes
@st.cache_resource
def load_models():
    return ProvencePruner(), SemanticHighlighter(), AgentSimulator()

pruner, highlighter, agent = load_models()

# Header
st.markdown('<h1 class="main-header">🧠 Agent Pruning Lab</h1>', unsafe_allow_html=True)
st.markdown('<p class="sub-header">Context Pruning com Provence + Semantic Highlighting para Agentes LLM</p>', unsafe_allow_html=True)

# Sidebar
with st.sidebar:
    st.markdown("### ⚙️ Configurações do Agente")
    
    model_choice = st.selectbox(
        "Modelo de Pruning",
        ["OpenProvence-XSmall (30M)", "OpenProvence-Small (110M)", "Simulação (Demo)"],
        index=2
    )
    
    pruning_threshold = st.slider(
        "Threshold de Pruning",
        0.0, 1.0, 0.5, 0.05,
        help="Sentenças com score abaixo são removidas"
    )
    
    highlight_mode = st.radio(
        "Modo de Visualização",
        ["Semantic Highlighting", "Binary Keep/Remove", "Heatmap de Relevância"],
        index=0
    )
    
    st.markdown("---")
    st.markdown("### 📚 Sobre")
    st.info("""
    **Provence** é um método de pruning de contexto desenvolvido pela Naver Labs Europe [^7^].
    
    **Semantic Highlighting** visualiza a relevância token-a-token do contexto [^11^].
    
    Combinação ideal para agentes que precisam processar documentos longos de forma eficiente.
    """)
    
    st.markdown("---")
    st.markdown("🔗 [Paper Provence](https://arxiv.org/abs/2501.16214)")
    st.markdown("🔗 [OpenProvence GitHub](https://github.com/hotchpotch/open_provence)")

# Tabs principais
tab1, tab2, tab3, tab4 = st.tabs([
    "🎮 Demo Interativo", 
    "📊 Benchmarks", 
    "🔬 Análise Técnica",
    "📖 Documentação"
])

with tab1:
    st.markdown("### 🎯 Simule um Agente de Pesquisa")
    
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.markdown("#### ❓ Query do Agente")
        query = st.text_area(
            "O que o agente precisa encontrar?",
            value="Quais são os principais benefícios do context pruning em sistemas RAG?",
            height=80
        )
        
        st.markdown("#### 📄 Documento de Entrada")
        document = st.text_area(
            "Cole um documento longo ou use o exemplo:",
            value="""A inteligência artificial tem revolucionado diversos setores da economia global. 
No campo da medicina, algoritmos de deep learning são utilizados para diagnóstico precoce de câncer. 
Em finanças, modelos preditivos analisam padrões de mercado. 
No entanto, um dos maiores desafios em sistemas de Retrieval-Augmented Generation (RAG) é o gerenciamento eficiente de contexto longo. 
Quando documentos irrelevantes são incluídos no prompt, o modelo desperdiça tokens e pode gerar respostas imprecisas. 
Técnicas como context pruning, especialmente métodos como Provence, permitem remover sentenças irrelevantes mantendo apenas o conteúdo necessário. 
Isso resulta em redução de latência e melhoria na qualidade das respostas. 
Estudos mostram que Provence consegue taxas de compressão de até 95% sem perda de performance [^1^]. 
Além disso, o Semantic Highlighting permite visualizar quais partes do texto são mais relevantes para uma query específica. 
Em sistemas multi-agente, onde múltiplos agentes trocam contexto, o pruning se torna ainda mais crítico para evitar propagação de ruído. 
A Naver Labs Europe desenvolveu o Provence como um modelo leve baseado em DeBERTa que realiza pruning e reranking em um único forward pass [^4^]. 
O modelo é treinado com labels de prata geradas por LLMs, identificando quais sentenças foram de fato utilizadas para gerar respostas. 
Resultados em benchmarks como Natural Questions e HotpotQA demonstram melhorias consistentes.""",
            height=300
        )
        
        run_button = st.button("🚀 Executar Pipeline do Agente", type="primary", use_container_width=True)
    
    with col2:
        st.markdown("#### 🧠 Simulação do Agente")
        
        if run_button:
            with st.spinner("🔍 Agente processando..."):
                # Simula o pipeline completo
                results = agent.run_pipeline(query, document, pruning_threshold)
                
                # Métricas
                col_m1, col_m2, col_m3 = st.columns(3)
                with col_m1:
                    st.markdown(f"""
                    <div class="metric-card">
                        <div class="metric-value">{results['compression_rate']}%</div>
                        <div class="metric-label">Compressão</div>
                    </div>
                    """, unsafe_allow_html=True)
                with col_m2:
                    st.markdown(f"""
                    <div class="metric-card">
                        <div class="metric-value">{results['tokens_saved']}</div>
                        <div class="metric-label">Tokens Economizados</div>
                    </div>
                    """, unsafe_allow_html=True)
                with col_m3:
                    st.markdown(f"""
                    <div class="metric-card">
                        <div class="metric-value">{results['relevance_score']:.2f}</div>
                        <div class="metric-label">Score de Relevância</div>
                    </div>
                    """, unsafe_allow_html=True)
                
                st.markdown("---")
                
                # Visualização do resultado
                st.markdown("#### ✂️ Resultado do Pruning")
                
                if highlight_mode == "Semantic Highlighting":
                    html_output = highlighter.render_semantic_highlighting(
                        results['sentences'], 
                        results['scores']
                    )
                    st.markdown(html_output, unsafe_allow_html=True)
                    
                elif highlight_mode == "Binary Keep/Remove":
                    for sent, score, kept in zip(results['sentences'], results['scores'], results['kept']):
                        css_class = "kept-text" if kept else "pruned-text"
                        st.markdown(f'<span class="{css_class}">{sent}</span> ', unsafe_allow_html=True)
                        
                else:  # Heatmap
                    fig = highlighter.render_heatmap(results['sentences'], results['scores'])
                    st.plotly_chart(fig, use_container_width=True)
                
                # Pensamento do agente
                st.markdown("---")
                st.markdown("#### 💭 Raciocínio do Agente")
                for thought in results['agent_thoughts']:
                    st.markdown(f'<div class="agent-thought">{thought}</div>', unsafe_allow_html=True)
        else:
            st.info("👈 Configure a query e clique em 'Executar Pipeline do Agente'")

with tab2:
    st.markdown("### 📊 Benchmarks de Performance")
    
    data = {
        "Método": ["Sem Pruning", "Truncation Simples", "Provence", "XProvence", "Semantic Highlighter"],
        "F1 Score": ["—", "42.3%", "66.76%", "58.97%", "46.37%"],
        "Compressão": ["0%", "60%", "80-95%", "75-85%", "70-80%"],
        "Latência Extra": ["0ms", "0ms", "~0ms*", "~0ms*", "5-10ms"],
        "Multilíngue": ["—", "Sim", "Inglês", "Sim", "Sim"]
    }
    
    import pandas as pd
    df = pd.DataFrame(data)
    st.dataframe(df, use_container_width=True, hide_index=True)
    
    st.caption("*Provence integra pruning com reranking, adicionando latência quase zero [^4^]")
    
    st.markdown("---")
    st.markdown("#### 🏆 Resultados em Datasets")
    
    col_b1, col_b2 = st.columns(2)
    with col_b1:
        st.markdown("**Natural Questions**")
        st.progress(0.89, text="Provence: 89% retenção de acurácia")
        st.progress(0.72, text="Baseline: 72% retenção de acurácia")
    with col_b2:
        st.markdown("**HotpotQA (Multi-hop)**")
        st.progress(0.85, text="Provence: 85% retenção de acurácia")
        st.progress(0.68, text="Baseline: 68% retenção de acurácia")

with tab3:
    st.markdown("### 🔬 Análise Técnica Detalhada")
    
    st.markdown("""
    #### Arquitetura Provence
    
    Provence utiliza uma arquitetura **dual-head** baseada em DeBERTa-V3 [^5^]:
    
    1. **Reranking Head**: Prediz score de relevância do documento inteiro
    2. **Pruning Head**: Gera máscara binária token-a-token indicando relevância
    
    Durante o treinamento, o modelo otimiza simultaneamente:
    - Binary Cross-Entropy para pruning
    - MSE para distilação do reranker original
    
    #### Semantic Highlighting vs Traditional Highlighting
    
    | Aspecto | Traditional (Keyword) | Semantic Highlighting |
    |---------|----------------------|----------------------|
    | Base | Matching literal | Similaridade semântica |
    | Sinônimos | ❌ Não detecta | ✅ Detecta |
    | Paráfrases | ❌ Não detecta | ✅ Detecta |
    | Custo | Baixo | Moderado |
    | Precisão | Baixa para RAG | Alta para RAG |
    
    #### Integração em Agentes
    
    Em workflows agentic, o pruning é aplicado em múltiplos pontos [^12^]:
    - **Pré-LLM**: Filtra documentos recuperados
    - **Inter-agente**: Resume contexto entre agentes
    - **Pós-tool**: Compacta outputs de ferramentas token-heavy
    """)

with tab4:
    st.markdown("### 📖 Como Usar Este Projeto")
    
    st.markdown("""
    #### 🚀 Deploy no Hugging Face Spaces
    
    1. Crie um novo Space em [huggingface.co/spaces](https://huggingface.co/spaces)
    2. Escolha **Streamlit** como SDK
    3. Faça push deste código:
    
    ```bash
    git clone https://huggingface.co/spaces/seu-username/agent-pruning-lab
    cd agent-pruning-lab
    # Copie os arquivos deste projeto
    git add .
    git commit -m "Initial commit"
    git push
    ```
    
    #### 🔧 Integração com Seu Agente
    
    ```python
    from components.provence_wrapper import ProvencePruner
    
    pruner = ProvencePruner(model_name="hotchpotch/open_provence_xsmall")
    
    # No seu pipeline RAG
    retrieved_docs = retriever.search(query)
    pruned_docs = pruner.prune(query, retrieved_docs, threshold=0.5)
    response = llm.generate(query, context=pruned_docs)
    ```
    
    #### 📚 Referências
    
    - [Provence Paper - ICLR 2025](https://arxiv.org/abs/2501.16214) [^7^]
    - [OpenProvence Implementation](https://github.com/hotchpotch/open_provence) [^8^]
    - [Semantic Highlighting Blog - Milvus](https://milvus.io/blog/semantic-highlighting-model-for-rag-context-pruning-and-token-saving.md) [^11^]
    - [Context Engineering - LangChain](https://www.langchain.com/blog/context-engineering-for-agents) [^12^]
    """)

# Footer
st.markdown("---")
st.markdown("""
<div style="text-align: center; color: #666; padding: 1rem;">
    <p>🧠 Agent Pruning Lab | Built with ❤️ for Hugging Face Spaces</p>
    <p>Demonstrando Provence + Semantic Highlighting para Agentes LLM Modernos</p>
</div>
""", unsafe_allow_html=True)