Update Gradio_UI.py

Gradio_UI.py

@@ -1,51 +1,191 @@
-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
-    """
+#!/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 file is None:
-        return gr.Textbox("No file uploaded", visible=True), file_uploads_log
+    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
+        self.file_upload_folder = file_upload_folder
+        if not os.path.exists(self.file_upload_folder):
+            os.makedirs(self.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",
+        ],
+    ):
+        import gradio as gr
 
-    try:
-        mime_type, _ = mimetypes.guess_type(file.name)
-    except Exception as e:
-        return gr.Textbox(f"Error: {e}", visible=True), file_uploads_log
+        if file is None:
+            return gr.Textbox("No file uploaded", visible=True), file_uploads_log
 
-    if mime_type not in allowed_file_types:
-        return gr.Textbox("File type disallowed", 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
 
-    # Sanitize file name
-    original_name = os.path.basename(file.name)
-    sanitized_name = re.sub(r"[^\w\-.]", "_", original_name)
+        if mime_type not in allowed_file_types:
+            return gr.Textbox("File type disallowed", visible=True), file_uploads_log
 
-    # Map mime type to extension
-    type_to_ext = {}
-    for ext, t in mimetypes.types_map.items():
-        if t not in type_to_ext:
-            type_to_ext[t] = ext
+        original_name = os.path.basename(file.name)
+        sanitized_name = re.sub(r"[^\w\-.]", "_", original_name)
 
-    name_without_ext = ".".join(sanitized_name.split(".")[:-1])
-    ext = type_to_ext.get(mime_type, "")
-    if not ext.startswith("."):
-        ext = "." + ext if ext else ""
-    sanitized_name = f"{name_without_ext}{ext}"
+        type_to_ext = {}
+        for ext, t in mimetypes.types_map.items():
+            if t not in type_to_ext:
+                type_to_ext[t] = ext
 
-    # Save file
-    file_path = os.path.join(self.file_upload_folder, sanitized_name)
-    with open(file_path, "wb") as f:
-        f.write(file.read())
+        name_without_ext = ".".join(sanitized_name.split(".")[:-1])
+        ext = type_to_ext.get(mime_type, "")
+        if not ext.startswith("."):
+            ext = "." + ext if ext else ""
+        sanitized_name = f"{name_without_ext}{ext}"
 
-    # Update log or return success message
-    file_uploads_log.append(file_path)
-    return gr.Textbox(f"File uploaded: {sanitized_name}", visible=True), file_uploads_log
+        file_path = os.path.join(self.file_upload_folder, sanitized_name)
+        with open(file_path, "wb") as f:
+            f.write(file.read())
 
+        file_uploads_log.append(file_path)
+        return gr.Textbox(f"File uploaded: {sanitized_name}", visible=True), file_uploads_log