Gradio_UI.py

#!/usr/bin/env python
# coding=utf-8
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import re
import json
from typing import Optional, List

import gradio as gr
from smolagents.agents import ActionStep, MultiStepAgent
from smolagents.memory import MemoryStep
from smolagents.utils import _is_package_available

def get_step_logs(step_log: MemoryStep):
    """
    Gera os logs intermediários do agente no formato [usuário, assistente].
    O usuário é 'None' pois são mensagens do sistema/assistente.
    """
    if isinstance(step_log, ActionStep):
        # Exibe o número do passo
        yield [None, f"**--- Passo {step_log.step_number} ---**"]
        
        # Exibe o pensamento/raciocínio do LLM
        if hasattr(step_log, "model_output") and step_log.model_output is not None:
            model_output = step_log.model_output.strip()
            model_output = re.sub(r"```\s*<end_code>", "```", model_output)
            yield [None, model_output]

        # Exibe as chamadas de ferramentas
        if hasattr(step_log, "tool_calls") and step_log.tool_calls is not None:
            tool_call = step_log.tool_calls[0]
            content = str(tool_call.arguments).strip()

            if tool_call.name == "python_interpreter":
                content = re.sub(r"\s*<end_code>\s*", "", content).strip()
                content = f"```python\n{content}\n```"

            yield [None, f"**🛠️ Ferramenta Usada: {tool_call.name}**\n{content}"]

            # Mostra os logs de execução ou erros
            if hasattr(step_log, "error") and step_log.error is not None:
                 yield [None, f"**💥 Erro na Execução:**\n```\n{str(step_log.error)}\n```"]
            elif hasattr(step_log, "observations") and step_log.observations and step_log.observations.strip():
                log_content = step_log.observations.strip()
                yield [None, f"**📝 Observação / Log de Execução:**\n```\n{log_content}\n```"]


def render_single_question(q_dict):
    """Renderiza um único dicionário de questão para uma tupla (texto, caminho_da_imagem)."""
    text = q_dict.get("enunciado", "")
    if "tabela_markdown" in q_dict: text += "\n\n" + q_dict["tabela_markdown"]
    if "alternativas" in q_dict and isinstance(q_dict["alternativas"], dict):
        text += "\n\n**Alternativas:**\n"
        for key, value in q_dict["alternativas"].items(): text += f"**{key})** {value}\n"
    if "resposta_correta" in q_dict: text += f"\n**Resposta Correta:** {q_dict['resposta_correta']}"
    if "resolucao_passo_a_passo" in q_dict: text += "\n\n**Resolução:**\n" + "\n".join(q_dict['resolucao_passo_a_passo'])
    
    img_key = q_dict.get("url_grafico") or q_dict.get("url_imagem")
    img_path = None
    if img_key:
        img_path = img_key.split(":")[-1].strip()
        if not os.path.exists(img_path): img_path = None
    return text, img_path
    
class GradioUI:
    """Interface para lançar o agente no Gradio."""
    def __init__(self, agent: MultiStepAgent):
        self.agent = agent

    def interact_with_agent(self, prompt, chat_history):
        """Função que gerencia a interação com o formato [pergunta, resposta]."""
        if not prompt.strip(): return chat_history, ""
        #Limpa a memória antes de cada execução para garantir novas respostas.
        self.agent.memory.reset()
        

        chat_history.append([prompt, ""])
        yield chat_history, "" # Limpa a caixa de texto

        step_log = None
        # Mostra os passos intermediários do agente
        for step_log in self.agent.run(task=prompt, stream=True):
            for user_msg, assistant_msg in get_step_logs(step_log):
                chat_history.append([user_msg, assistant_msg])
                yield chat_history, ""

        # Renderiza a resposta final
        if step_log:
            
            final_answer = step_log.final_answer if hasattr(step_log, 'final_answer') else step_log
            
            # Remove o log do último passo para não poluir a resposta final
            chat_history.pop() 
            chat_history.append([None, "--- \n\n ### Resposta Final"])

            if isinstance(final_answer, list):
                # Se for uma lista, processa e exibe cada item individualmente
                for i, item in enumerate(final_answer):
                    if isinstance(item, dict):
                        q_text, q_image_path = render_single_question(item)
                        chat_history.append([None, f"#### Questão {i+1} de {len(final_answer)}\n{q_text}"])
                        if q_image_path:
                            chat_history.append([None, (q_image_path,)])
                yield chat_history, ""
            elif isinstance(final_answer, dict):
                # Se for um item único
                q_text, q_image_path = render_single_question(final_answer)
                chat_history.append([None, q_text])
                if q_image_path:
                    chat_history.append([None, (q_image_path,)])
                yield chat_history, ""
            else:
                chat_history.append([None, f"**Resposta Final:** {str(final_answer)}"])
                yield chat_history, ""
            
        else:
            chat_history.append([None, "**ERRO:** O agente não produziu nenhuma saída."])
            yield chat_history, ""


    def launch(self, **kwargs):
        with gr.Blocks(fill_height=True) as demo:
            gr.Markdown("# Agente Gerador de Questões de Matemática")

            
            #chatbot = gr.Chatbot(label="Interação com o Agente", bubble_full_width=True, height=700)
            chatbot = gr.Chatbot(
                label="Interação com o Agente", 
                bubble_full_width=True, 
                height=700,
                #instrui o chatbot a procurar por \(...\) e renderizar como matemática
                latex_delimiters=[{"left": "\\(", "right": "\\)", "display": False}]
            )

            
            text_input = gr.Textbox(lines=1, label="Sua Tarefa", placeholder="Ex: habilidade_bncc:(EM13MAT501) Investigar ..., nivel_dificuldade: médio (fácil ou difícil), quantidade:número, tipo_questao: multipla_escolha ou dissertativa, observação: aqui é um campo facultativo")
            text_input.submit(self.interact_with_agent, [text_input, chatbot], [chatbot, text_input])
        demo.launch(debug=True, **kwargs)

__all__ = ["GradioUI"]