Update Gradio_UI.py

Gradio_UI.py

@@ -1,135 +1,3 @@
-#!/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
-from typing import Optional
-
-from smolagents.agent_types import AgentAudio, AgentImage, AgentText, handle_agent_output_types
-from smolagents.agents import ActionStep
-from smolagents.memory import MemoryStep
-from smolagents.utils import _is_package_available
-
-
-def pull_messages_from_step(step_log: MemoryStep):
-    import gradio as gr
-
-    if isinstance(step_log, ActionStep):
-        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}**")
-
-        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)
-            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)
-
-        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)}"
-            args = first_tool_call.arguments
-            content = str(args.get("answer", str(args))) if isinstance(args, dict) else str(args).strip()
-
-            if used_code:
-                content = re.sub(r"```.*?\n", "", content)
-                content = re.sub(r"\s*<end_code>\s*", "", 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
-
-            if hasattr(step_log, "observations") and step_log.observations and step_log.observations.strip():
-                log_content = step_log.observations.strip()
-                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"},
-                )
-
-            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"
-
-        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"})
-
-        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:,}"
-            )
-            step_footnote += token_str
-        if hasattr(step_log, "duration"):
-            duration = round(float(step_log.duration), 2)
-            step_footnote += f" | Duration: {duration}"
-
-        step_footnote = f"""<span style="color: #bbbbc2; font-size: 12px;">{step_footnote}</span> """
-        yield gr.ChatMessage(role="assistant", content=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):
-    if not _is_package_available("gradio"):
-        raise ModuleNotFoundError(
-            "Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`"
-        )
-
-    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):
-        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 = handle_agent_output_types(step_log)
-
-    import gradio as gr
-    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:
     def __init__(self, agent, file_upload_folder: str = "./uploads"):
         self.agent = agent
@@ -189,3 +57,24 @@ class GradioUI:
 
         file_uploads_log.append(file_path)
         return gr.Textbox(f"File uploaded: {sanitized_name}", visible=True), file_uploads_log
+
+    def build_ui(self):
+        import gradio as gr
+
+        chatbot = gr.Chatbot()
+        msg = gr.Textbox(placeholder="Ask something...", label="Your message")
+        clear = gr.Button("Clear")
+        file_uploads_log = []
+
+        with gr.Blocks() as demo:
+            with gr.Row():
+                with gr.Column():
+                    chatbot.render()
+                    msg.render()
+                    clear.render()
+
+            msg.submit(self.interact_with_agent, [msg, chatbot], chatbot)
+            clear.click(lambda: None, None, chatbot, queue=False)
+
+        return demo
+