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 mimetypes
import os
import re
import shutil
from typing import Optional

from smolagents.agent_types import AgentAudio, AgentImage, AgentText, handle_agent_output_types
from smolagents.agents import ActionStep, MultiStepAgent
from smolagents.memory import MemoryStep
from smolagents.utils import _is_package_available

def pull_messages_from_step(step_log: MemoryStep):
    """
    Extract ChatMessage objects from agent steps with proper nesting.
    This is where we transform the agent's step-by-step reasoning,
    tool calls, and logs into user-friendly gradio ChatMessage objects.
    """
    import gradio as gr

    if isinstance(step_log, ActionStep):
        # Output the step number
        step_number = f"Step {step_log.step_number}" if step_log.step_number is not None else ""
        yield gr.ChatMessage(role="assistant", content=f"**{step_number}**")

        # First yield the thought/reasoning from the LLM
        if hasattr(step_log, "model_output") and step_log.model_output is not None:
            # Clean up the LLM output
            model_output = step_log.model_output.strip()
            # Remove any trailing <end_code> and extra backticks
            model_output = re.sub(r"```\s*<end_code>", "```", model_output)
            model_output = re.sub(r"<end_code>\s*```", "```", model_output)
            model_output = re.sub(r"```\s*\n\s*<end_code>", "```", model_output)
            model_output = model_output.strip()

            yield gr.ChatMessage(role="assistant", content=model_output)

        # For tool calls
        if hasattr(step_log, "tool_calls") and step_log.tool_calls is not None:
            first_tool_call = step_log.tool_calls[0]
            used_code = first_tool_call.name == "python_interpreter"
            parent_id = f"call_{len(step_log.tool_calls)}"

            # Display the arguments used for the tool call
            args = first_tool_call.arguments
            if isinstance(args, dict):
                content = str(args.get("answer", str(args)))
            else:
                content = str(args).strip()

            if used_code:
                # Clean up content by removing code blocks
                content = re.sub(r"```.*?\n", "", content)
                content = re.sub(r"\s*<end_code>\s*", "", content)
                content = content.strip()
                if not content.startswith("```python"):
                    content = f"```python\n{content}\n```"

            parent_message_tool = gr.ChatMessage(
                role="assistant",
                content=content,
                metadata={
                    "title": f"🛠️ Used tool {first_tool_call.name}",
                    "id": parent_id,
                    "status": "pending",
                },
            )
            yield parent_message_tool

            # Observations or logs from the tool call
            if hasattr(step_log, "observations") and step_log.observations is not None and step_log.observations.strip():
                log_content = step_log.observations.strip()
                if log_content:
                    log_content = re.sub(r"^Execution logs:\s*", "", log_content)
                    yield gr.ChatMessage(
                        role="assistant",
                        content=log_content,
                        metadata={
                            "title": "📝 Execution Logs",
                            "parent_id": parent_id,
                            "status": "done",
                        },
                    )

            # Handle any errors
            if hasattr(step_log, "error") and step_log.error is not None:
                yield gr.ChatMessage(
                    role="assistant",
                    content=str(step_log.error),
                    metadata={"title": "💥 Error", "parent_id": parent_id, "status": "done"},
                )

            parent_message_tool.metadata["status"] = "done"

        # Standalone errors
        elif hasattr(step_log, "error") and step_log.error is not None:
            yield gr.ChatMessage(role="assistant", content=str(step_log.error), metadata={"title": "💥 Error"})

        # Token counts, durations, etc.
        step_footnote = f"{step_number}"
        if hasattr(step_log, "input_token_count") and hasattr(step_log, "output_token_count"):
            token_str = (
                f" | Input-tokens:{step_log.input_token_count:,} "
                f"| Output-tokens:{step_log.output_token_count:,}"
            )
            step_footnote += token_str
        if hasattr(step_log, "duration"):
            step_duration = (
                f" | Duration: {round(float(step_log.duration), 2)}"
                if step_log.duration
                else None
            )
            step_footnote += step_duration if step_duration else ""
        step_footnote_html = f"""<span style="color: #bbbbc2; font-size: 12px;">{step_footnote}</span> """
        yield gr.ChatMessage(role="assistant", content=f"{step_footnote_html}")
        yield gr.ChatMessage(role="assistant", content="-----")


def stream_to_gradio(agent, task: str, reset_agent_memory: bool = False, additional_args: Optional[dict] = None):
    """
    Runs an agent with the given task and streams the messages as gradio ChatMessages.
    """
    if not _is_package_available("gradio"):
        raise ModuleNotFoundError(
            "Please install 'gradio' extra to use Gradio: `pip install 'smolagents[gradio]'`"
        )
    import gradio as gr

    total_input_tokens = 0
    total_output_tokens = 0

    # Run the agent in streaming mode
    for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
        # If the model tracks token usage:
        if hasattr(agent.model, "last_input_token_count"):
            total_input_tokens += agent.model.last_input_token_count
            total_output_tokens += agent.model.last_output_token_count
            if isinstance(step_log, ActionStep):
                step_log.input_token_count = agent.model.last_input_token_count
                step_log.output_token_count = agent.model.last_output_token_count

        # Convert each step into user-friendly messages
        for message in pull_messages_from_step(step_log):
            yield message

    # The last step_log is presumably the final answer
    final_answer = step_log
    final_answer = handle_agent_output_types(final_answer)

    # Convert the final_answer into a string for the user
    final_answer_str = getattr(final_answer, "final_answer", "")
    final_answer_str = f"\n{final_answer_str}\n"

    # Yield one last message containing the final answer
    yield gr.ChatMessage(role="assistant", content=final_answer_str)

    return final_answer_str


class GradioUI:
    """
    A one-line interface to launch your agent in Gradio.
    Features:
    - Chatbot panel for user messages and step-by-step agent reasoning
    - 'Final Answer' section that clearly shows the final result
    - (Optional) file upload for extra data the agent might use
    """

    def __init__(self, agent: MultiStepAgent, file_upload_folder: Optional[str] = None):
        if not _is_package_available("gradio"):
            raise ModuleNotFoundError(
                "Please install 'gradio' extra to use Gradio: `pip install 'smolagents[gradio]'`"
            )
        self.agent = agent
        self.file_upload_folder = file_upload_folder
        if self.file_upload_folder is not None and not os.path.exists(file_upload_folder):
            os.mkdir(file_upload_folder)

    def interact_with_agent(self, prompt, messages, final_answer_state):
        """
        This function is called whenever the user submits a new query.
        We append the user's message to the chat, then stream the agent's steps
        back to the chatbot widget, and finally store the final answer.
        """
        import gradio as gr

        # Add the user's new message to the conversation
        messages.append(gr.ChatMessage(role="user", content=prompt))
        yield messages, final_answer_state

        # Stream out each step of the agent's process
        for msg in stream_to_gradio(self.agent, task=prompt, reset_agent_memory=False):
            messages.append(msg)
            yield messages, final_answer_state

        # Update the final answer state once the agent is done
        final_answer_state = msg.content if isinstance(msg, gr.ChatMessage) else ""
        yield messages, final_answer_state

    def upload_file(
            self,
            file,
            file_uploads_log,
            allowed_file_types=[
                "application/pdf",
                "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                "text/plain",
            ],
    ):
        """
        Handle file uploads, default allowed types are .pdf, .docx, and .txt
        """
        import gradio as gr

        if file is None:
            return gr.Textbox("No file uploaded", visible=True), file_uploads_log

        # Attempt to detect mime type
        try:
            mime_type, _ = mimetypes.guess_type(file.name)
        except Exception as e:
            return gr.Textbox(f"Error: {e}", visible=True), file_uploads_log

        # Check if file is allowed
        if mime_type not in allowed_file_types:
            return gr.Textbox("File type disallowed", visible=True), file_uploads_log

        # Sanitize and rename
        original_name = os.path.basename(file.name)
        sanitized_name = re.sub(r"[^\w\-.]", "_", original_name)

        type_to_ext = {}
        for ext, t in mimetypes.types_map.items():
            if t not in type_to_ext:
                type_to_ext[t] = ext

        # Append the correct extension for the mime type
        base_name = ".".join(sanitized_name.split(".")[:-1])
        extension = type_to_ext.get(mime_type, "")
        final_name = f"{base_name}{extension}".strip()

        # Save the file
        file_path = os.path.join(self.file_upload_folder, os.path.basename(final_name))
        shutil.copy(file.name, file_path)

        return gr.Textbox(f"File uploaded: {file_path}", visible=True), file_uploads_log + [file_path]

    def log_user_message(self, text_input, file_uploads_log):
        """
        Combines the user's text input with any file references.
        We pass this along to the agent so it knows about available files.
        """
        combined_input = text_input
        if file_uploads_log:
            combined_input += f"\nYou have been provided these files: {file_uploads_log}"
        return combined_input, ""

    def launch(self, **kwargs):
        """
        Build and launch the Gradio Blocks interface with:
          - A Markdown introduction
          - A chat panel
          - A file upload option (optional)
          - A final answer panel
          - Example usage instructions
        """
        import gradio as gr

        with gr.Blocks() as demo:
            # Heading and instructions
            gr.Markdown("""
            # 😎 Pink Glasses Agent ☀️  
            A cheerful AI that filters out negativity and shares only uplifting, feel-good responses.  
            
            Ask anything — the agent thinks step by step and delivers a happy final answer. It can also fetch the latest news using the `fetch_news` tool powered by DuckDuckGo.

            ---
            """)

            with gr.Row():
                with gr.Column():
                    stored_messages = gr.State([])
                    file_uploads_log = gr.State([])
                    final_answer_state = gr.State("")

                    chatbot = gr.Chatbot(
                        label="Pink Glasses Agent",
                        type="messages",
                        avatar_images=(None, "https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png"),
                        height=500,
                    )

                    # Optional file upload
                    if self.file_upload_folder is not None:
                        upload_file = gr.File(label="Upload a file")
                        upload_status = gr.Textbox(label="Upload Status", interactive=False, visible=False)
                        upload_file.change(
                            self.upload_file,
                            [upload_file, file_uploads_log],
                            [upload_status, file_uploads_log],
                        )

                    text_input = gr.Textbox(lines=1, label="Your Input")
                    text_input.submit(
                        self.log_user_message,
                        [text_input, file_uploads_log],
                        [stored_messages, text_input],
                    ).then(
                        self.interact_with_agent,
                        [stored_messages, chatbot, final_answer_state],
                        [chatbot, final_answer_state],
                    )

                with gr.Column():
                    final_answer_display = gr.Markdown("## Final Answer")
                    final_answer_state.change(
                        lambda state: f"## Final Answer\n\n{state}",
                        inputs=final_answer_state,
                        outputs=final_answer_display,
                    )

            gr.Markdown("""
            ---
            # Example Usage
            - **Ask about the latest news**: 
              > "What's going on in the world right now?"
            - **Use Tools**: 
              The agent can fetch the latest news using DuckDuckGo.
            """)

        # Optional: share=True if you want a public link
        demo.launch(debug=True, share=True, **kwargs)