Update Gradio_UI.py

Gradio_UI.py

@@ -1,102 +1,296 @@
-import gradio as gr
-from smolagents.agent_types import AgentImage, AgentText
-import markdown
+#!/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
 
-CSS = """
-.contain { display: flex; flex-direction: column; }
-.gr-chatbot { height: 75vh; }
-#component-0 { height: 100%; }
-"""
+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
 
-class GradioUI:
-    def __init__(self, agent):
-        self.agent = agent
-        with gr.Blocks(css=CSS) as self.interface:
-            gr.Markdown("# smolagents")
-            
-            # Use a two-column layout to place the chatbot and image side-by-side
-            with gr.Row():
-                with gr.Column(scale=2):
-                    chatbot = gr.Chatbot(height=600)
-                with gr.Column(scale=1):
-                    # Dedicated component for image output
-                    image_output = gr.Image(label="Generated Image", type="pil", height=600)
-            
-            # Input textbox and clear button
-            with gr.Row():
-                msg = gr.Textbox(
-                    show_label=False,
-                    placeholder="Enter text and press enter",
-                    container=False
+
+def pull_messages_from_step(
+    step_log: MemoryStep,
+):
+    """Extract ChatMessage objects from agent steps with proper nesting"""
+    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, handling multiple possible formats
+            model_output = re.sub(r"```\s*<end_code>", "```", model_output)  # handles ```<end_code>
+            model_output = re.sub(r"<end_code>\s*```", "```", model_output)  # handles <end_code>```
+            model_output = re.sub(r"```\s*\n\s*<end_code>", "```", model_output)  # handles ```\n<end_code>
+            model_output = model_output.strip()
+            yield gr.ChatMessage(role="assistant", content=model_output)
+
+        # For tool calls, create a parent message
+        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)}"
+
+            # Tool call becomes the parent message with timing info
+            # First we will handle arguments based on type
+            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 the content by removing any end code tags
+                content = re.sub(r"```.*?\n", "", content)  # Remove existing code blocks
+                content = re.sub(r"\s*<end_code>\s*", "", content)  # Remove end_code tags
+                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
+
+            # Nesting execution logs under the tool call if they exist
+            if hasattr(step_log, "observations") and (
+                step_log.observations is not None and step_log.observations.strip()
+            ):  # Only yield execution logs if there's actual content
+                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=f"{log_content}",
+                        metadata={"title": "📝 Execution Logs", "parent_id": parent_id, "status": "done"},
+                    )
+
+            # Nesting any errors under the tool call
+            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"},
                 )
-                clear = gr.ClearButton([msg, chatbot, image_output])
-
-            self.interface.queue()
-            
-            # The submit function now has the image component as an output
-            msg.submit(
-                self.interact_with_agent,
-                [msg, chatbot],
-                [msg, chatbot, image_output],
+
+            # Update parent message metadata to done status without yielding a new message
+            parent_message_tool.metadata["status"] = "done"
+
+        # Handle standalone errors but not from tool calls
+        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"})
+
+        # Calculate duration and token information
+        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:,} | Output-tokens:{step_log.output_token_count:,}"
             )
-            
-    def interact_with_agent(self, text, history):
-        history = history or []
-        yield text, history, None # Clear the image output on new input
-        
-        history.append([text, None])
-        yield "", history, None
-
-        # This part of the code now correctly handles both text and image outputs
-        output_stream = self.agent.run(text)
-        
-        last_yielded_image = None
-        
-        for item in output_stream:
-            if isinstance(item, AgentText):
-                # Process and update chatbot for text output
-                text_to_yield = markdown.markdown(item.content)
-                history[-1][1] = text_to_yield
-            elif isinstance(item, AgentImage):
-                # Update the image variable for image output
-                last_yielded_image = item.image
-            
-            # Yield both the chatbot history and the current image
-            yield "", history, last_yielded_image
-            
-    def launch(self):
-        self.interface.launch()
-
-if __name__ == '__main__':
-    from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool
-    import yaml
-    
-    # Example agent setup from your original app.py
-    # This part should be in your main app.py, not here
-    # but is included for a full, runnable example
-    final_answer = load_tool("tools.final_answer.FinalAnswerTool")
-    web_search = DuckDuckGoSearchTool()
-    image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
-    
-    model = HfApiModel(
-        max_tokens=2096,
-        temperature=0.5,
-        model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
-        custom_role_conversions=None,
-    )
-    
-    with open("prompts.yaml", 'r') as stream:
-        prompt_templates = yaml.safe_load(stream)
-        agent = CodeAgent(
-            model=model,
-            tools=[final_answer, web_search, image_generation_tool],
-            max_steps=6,
-            verbosity_level=1,
-            grammar=None,
-            planning_interval=None,
-            name=None,
-            description=None,
-            prompt_templates=prompt_templates
+            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
+        step_footnote = f"""<span style="color: #bbbbc2; font-size: 12px;">{step_footnote}</span> """
+        yield gr.ChatMessage(role="assistant", content=f"{step_footnote}")
+        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 from the agent as gradio ChatMessages."""
+    if not _is_package_available("gradio"):
+        raise ModuleNotFoundError(
+            "Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`"
         )
-    
-    GradioUI(agent).launch()
+    import gradio as gr
+
+    total_input_tokens = 0
+    total_output_tokens = 0
+
+    for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
+        # Track tokens if model provides them
+        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
+
+        for message in pull_messages_from_step(
+            step_log,
+        ):
+            yield message
+
+    final_answer = step_log  # Last log is the run's final_answer
+    final_answer = handle_agent_output_types(final_answer)
+
+    if isinstance(final_answer, AgentText):
+        yield gr.ChatMessage(
+            role="assistant",
+            content=f"**Final answer:**\n{final_answer.to_string()}\n",
+        )
+    elif isinstance(final_answer, AgentImage):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"path": final_answer.to_string(), "mime_type": "image/png"},
+        )
+    elif isinstance(final_answer, AgentAudio):
+        yield gr.ChatMessage(
+            role="assistant",
+            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
+        )
+    else:
+        yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
+
+
+class GradioUI:
+    """A one-line interface to launch your agent in Gradio"""
+
+    def __init__(self, agent: MultiStepAgent, file_upload_folder: str | None = None):
+        if not _is_package_available("gradio"):
+            raise ModuleNotFoundError(
+                "Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`"
+            )
+        self.agent = agent
+        self.file_upload_folder = file_upload_folder
+        if self.file_upload_folder is not None:
+            if not os.path.exists(file_upload_folder):
+                os.mkdir(file_upload_folder)
+
+    def interact_with_agent(self, prompt, messages):
+        import gradio as gr
+
+        messages.append(gr.ChatMessage(role="user", content=prompt))
+        yield messages
+        for msg in stream_to_gradio(self.agent, task=prompt, reset_agent_memory=False):
+            messages.append(msg)
+            yield messages
+        yield messages
+
+    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
+
+        try:
+            mime_type, _ = mimetypes.guess_type(file.name)
+        except Exception as e:
+            return gr.Textbox(f"Error: {e}", visible=True), file_uploads_log
+
+        if mime_type not in allowed_file_types:
+            return gr.Textbox("File type disallowed", visible=True), file_uploads_log
+
+        # Sanitize file name
+        original_name = os.path.basename(file.name)
+        sanitized_name = re.sub(
+            r"[^\w\-.]", "_", original_name
+        )  # Replace any non-alphanumeric, non-dash, or non-dot characters with underscores
+
+        type_to_ext = {}
+        for ext, t in mimetypes.types_map.items():
+            if t not in type_to_ext:
+                type_to_ext[t] = ext
+
+        # Ensure the extension correlates to the mime type
+        sanitized_name = sanitized_name.split(".")[:-1]
+        sanitized_name.append("" + type_to_ext[mime_type])
+        sanitized_name = "".join(sanitized_name)
+
+        # Save the uploaded file to the specified folder
+        file_path = os.path.join(self.file_upload_folder, os.path.basename(sanitized_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):
+        return (
+            text_input
+            + (
+                f"\nYou have been provided with these files, which might be helpful or not: {file_uploads_log}"
+                if len(file_uploads_log) > 0
+                else ""
+            ),
+            "",
+        )
+
+    def launch(self, **kwargs):
+        import gradio as gr
+
+        with gr.Blocks(fill_height=True) as demo:
+            stored_messages = gr.State([])
+            file_uploads_log = gr.State([])
+            chatbot = gr.Chatbot(
+                label="Agent",
+                type="messages",
+                avatar_images=(
+                    None,
+                    "https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png",
+                ),
+                resizeable=True,
+                scale=1,
+            )
+            # If an upload folder is provided, enable the upload feature
+            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="Chat Message")
+            text_input.submit(
+                self.log_user_message,
+                [text_input, file_uploads_log],
+                [stored_messages, text_input],
+            ).then(self.interact_with_agent, [stored_messages, chatbot], [chatbot])
+
+        demo.launch(debug=True, share=True, **kwargs)
+
+
+__all__ = ["stream_to_gradio", "GradioUI"]