ui/playground_tab.py

"""
Aba de Playground de Parâmetros
Permite experimentar com diferentes configurações e comparar resultados
"""
import time
import gradio as gr
from src.database import DatabaseManager
from src.embeddings import EmbeddingManager
from src.generation import GenerationManager


def create_playground_tab(
    db_manager: DatabaseManager,
    embedding_manager: EmbeddingManager,
    generation_manager: GenerationManager,
    session_id: str
):
    """Cria aba de playground de parâmetros"""

    with gr.Tab(" Playground de Parâmetros"):
        gr.Markdown("""
        ## Experimente com Diferentes Parâmetros

        Compare como diferentes configurações afetam os resultados:
        - **Top K**: Quantos chunks recuperar
        - **Chunk Size**: Tamanho dos blocos de texto
        - **Temperature**: Criatividade da resposta
        - **Max Tokens**: Tamanho máximo da resposta
        """)

        with gr.Row():
            test_query = gr.Textbox(
                label=" Query de Teste",
                placeholder="Digite uma pergunta para testar com diferentes parâmetros...",
                lines=2
            )

        gr.Markdown("### Configuração A vs Configuração B")

        with gr.Row():
            # Configuração A
            with gr.Column():
                gr.Markdown("####  Configuração A")

                top_k_a = gr.Slider(1, 20, value=3, step=1, label="Top K")
                temperature_a = gr.Slider(0.0, 2.0, value=0.1, step=0.1, label="Temperature")
                max_tokens_a = gr.Slider(50, 2048, value=256, step=50, label="Max Tokens")

                run_a_btn = gr.Button("  Executar A", variant="primary")

                gr.Markdown("**Resultado A**")
                result_a = gr.Textbox(label="Resposta", lines=10, interactive=False)

                with gr.Accordion(" Detalhes A", open=False):
                    contexts_a = gr.Dataframe(
                        headers=["Rank", "Score", "Fonte"],
                        label="Contextos Recuperados"
                    )
                    metrics_a = gr.JSON(label="Métricas")

            # Configuração B
            with gr.Column():
                gr.Markdown("####  Configuração B")

                top_k_b = gr.Slider(1, 20, value=8, step=1, label="Top K")
                temperature_b = gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="Temperature")
                max_tokens_b = gr.Slider(50, 2048, value=512, step=50, label="Max Tokens")

                run_b_btn = gr.Button("  Executar B", variant="primary")

                gr.Markdown("**Resultado B**")
                result_b = gr.Textbox(label="Resposta", lines=10, interactive=False)

                with gr.Accordion(" Detalhes B", open=False):
                    contexts_b = gr.Dataframe(
                        headers=["Rank", "Score", "Fonte"],
                        label="Contextos Recuperados"
                    )
                    metrics_b = gr.JSON(label="Métricas")

        with gr.Row():
            compare_btn = gr.Button("  Comparar Resultados Lado a Lado", size="lg")

        with gr.Row():
            comparison_output = gr.Markdown("Clique em 'Comparar Resultados' após executar ambas configurações")

        # Função para executar configuração
        def run_config(query, top_k, temperature, max_tokens):
            if not query or not query.strip():
                return "  Digite uma query", [], {}

            start_time = time.time()

            # Retrieve
            retrieve_start = time.time()
            query_embedding = embedding_manager.encode_single(query, normalize=True)
            contexts = db_manager.search_similar(query_embedding, k=int(top_k), session_id=session_id)
            retrieve_time = (time.time() - retrieve_start) * 1000

            # Contextos para tabela
            contexts_table = []
            for i, ctx in enumerate(contexts, 1):
                contexts_table.append([
                    i,
                    f"{ctx['score']:.4f}",
                    ctx['title']
                ])

            # Generate
            generate_start = time.time()
            prompt = generation_manager.build_rag_prompt(query, contexts)
            response = generation_manager.generate(
                prompt,
                temperature=float(temperature),
                max_tokens=int(max_tokens)
            )
            generate_time = (time.time() - generate_start) * 1000

            total_time = (time.time() - start_time) * 1000

            # Métricas
            metrics = {
                "top_k": int(top_k),
                "temperature": float(temperature),
                "max_tokens": int(max_tokens),
                "num_contexts": len(contexts),
                "retrieval_time_ms": retrieve_time,
                "generation_time_ms": generate_time,
                "total_time_ms": total_time,
                "avg_context_score": sum(c['score'] for c in contexts) / len(contexts) if contexts else 0,
                "response_length": len(response)
            }

            return response, contexts_table, metrics

        # Função de comparação
        def compare_results(result_a_text, metrics_a_data, result_b_text, metrics_b_data):
            if not result_a_text or not result_b_text:
                return "  Execute ambas configurações antes de comparar"

            if not metrics_a_data or not metrics_b_data:
                return "  Métricas não disponíveis"

            comparison_md = f"""
            ##  Comparação de Resultados

            ###  Performance
            | Métrica | Config A | Config B | Diferença |
            |---------|----------|----------|-----------|
            | Tempo Total | {metrics_a_data.get('total_time_ms', 0):.0f}ms | {metrics_b_data.get('total_time_ms', 0):.0f}ms | {abs(metrics_a_data.get('total_time_ms', 0) - metrics_b_data.get('total_time_ms', 0)):.0f}ms |
            | Retrieval | {metrics_a_data.get('retrieval_time_ms', 0):.0f}ms | {metrics_b_data.get('retrieval_time_ms', 0):.0f}ms | {abs(metrics_a_data.get('retrieval_time_ms', 0) - metrics_b_data.get('retrieval_time_ms', 0)):.0f}ms |
            | Geração | {metrics_a_data.get('generation_time_ms', 0):.0f}ms | {metrics_b_data.get('generation_time_ms', 0):.0f}ms | {abs(metrics_a_data.get('generation_time_ms', 0) - metrics_b_data.get('generation_time_ms', 0)):.0f}ms |

            ### Qualidade dos Contextos
            | Métrica | Config A | Config B |
            |---------|----------|----------|
            | Top K | {metrics_a_data.get('top_k', 0)} | {metrics_b_data.get('top_k', 0)} |
            | Contextos Retornados | {metrics_a_data.get('num_contexts', 0)} | {metrics_b_data.get('num_contexts', 0)} |
            | Score Médio | {metrics_a_data.get('avg_context_score', 0):.4f} | {metrics_b_data.get('avg_context_score', 0):.4f} |

            ### Resposta
            | Métrica | Config A | Config B |
            |---------|----------|----------|
            | Temperature | {metrics_a_data.get('temperature', 0)} | {metrics_b_data.get('temperature', 0)} |
            | Max Tokens | {metrics_a_data.get('max_tokens', 0)} | {metrics_b_data.get('max_tokens', 0)} |
            | Tamanho da Resposta | {metrics_a_data.get('response_length', 0)} chars | {metrics_b_data.get('response_length', 0)} chars |

            ### Insights
            - **Top K**: Config {'A' if metrics_a_data.get('top_k', 0) > metrics_b_data.get('top_k', 0) else 'B'} recuperou mais contextos
            - **Performance**: Config {'A' if metrics_a_data.get('total_time_ms', 0) < metrics_b_data.get('total_time_ms', 0) else 'B'} foi mais rápida
            - **Criatividade**: Config {'A' if metrics_a_data.get('temperature', 0) > metrics_b_data.get('temperature', 0) else 'B'} tem maior temperature (mais criativa)
            - **Tamanho**: Config {'A' if metrics_a_data.get('response_length', 0) > metrics_b_data.get('response_length', 0) else 'B'} gerou resposta mais longa
            """

            return comparison_md

        # Estado para armazenar resultados
        result_a_state = gr.State("")
        result_b_state = gr.State("")
        metrics_a_state = gr.State({})
        metrics_b_state = gr.State({})

        # Conecta eventos Config A
        run_a_btn.click(
            fn=run_config,
            inputs=[test_query, top_k_a, temperature_a, max_tokens_a],
            outputs=[result_a, contexts_a, metrics_a]
        ).then(
            fn=lambda r, m: (r, m),
            inputs=[result_a, metrics_a],
            outputs=[result_a_state, metrics_a_state]
        )

        # Conecta eventos Config B
        run_b_btn.click(
            fn=run_config,
            inputs=[test_query, top_k_b, temperature_b, max_tokens_b],
            outputs=[result_b, contexts_b, metrics_b]
        ).then(
            fn=lambda r, m: (r, m),
            inputs=[result_b, metrics_b],
            outputs=[result_b_state, metrics_b_state]
        )

        # Conecta comparação
        compare_btn.click(
            fn=compare_results,
            inputs=[result_a_state, metrics_a_state, result_b_state, metrics_b_state],
            outputs=[comparison_output]
        )

    return {
        "test_query": test_query
    }