Update app UI

gradio_ui.py

@@ -13,21 +13,101 @@
 # 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 shutil
+import base64
 from typing import Optional
 
+import gradio as gr
+
 from smolagents.agent_types import AgentAudio, AgentImage, AgentText
 from smolagents.agents import MultiStepAgent, PlanningStep
 from smolagents.memory import ActionStep, FinalAnswerStep, MemoryStep
 from smolagents.utils import _is_package_available
-import base64
+
+CUSTOM_CSS = """
+    .gradio-container {min-height: 100vh;} 
+    .content-wrap {padding-bottom: 60px;}
+    .full-width-btn {
+        width: 100% !important;
+        height: 50px !important;
+        font-size: 18px !important;
+        margin-top: 20px !important;
+        background: linear-gradient(45deg, #FF6B6B, #4ECDC4) !important;
+        color: white !important;
+        border: none !important;
+    }
+    .full-width-btn:hover {
+        background: linear-gradient(45deg, #FF5252, #3CB4AC) !important;
+    }
+    """
+
 
 def image_to_base64(image_path):
     with open(image_path, "rb") as image_file:
         return base64.b64encode(image_file.read()).decode('utf-8')
 
+
+def create_header():
+    with gr.Row():
+        with gr.Column(scale=1):
+            if os.path.exists("static/aivn_logo.png"):
+                logo_base64 = image_to_base64("static/aivn_logo.png")
+                gr.HTML(f"""
+                <img src="data:image/png;base64,{logo_base64}" 
+                    alt="Logo" 
+                    style="height: 120px; width: auto; margin-right: 20px; margin-bottom: 20px;">
+                """)
+            else:
+                gr.HTML("""
+                <div style="height: 120px; display: flex; align-items: center; justify-content: center; font-size: 24px; font-weight: bold;">
+                    AI VIETNAM
+                </div>
+                """)
+        with gr.Column(scale=4):
+            gr.Markdown(
+                """
+                <div style="display: flex; justify-content: space-between; align-items: center; padding: 0 15px;">
+                    <div>
+                        <h1 style="margin-bottom:0;">📰 Smart News Agent</h1>
+                        <p style="margin-top: 0.5em; color: #666;">🚀 AIO2024 Module 10 🤗</p>
+                        <p style="margin-top: 0.5em; color: #2c3e50;">🗞️ Real-time News Fetch & Summarization</p>
+                        <p style="margin-top: 0.2em; color: #7f8c8d;">🔍 Topic Classification & Insight Extraction</p>
+                    </div>
+                </div>
+                """)
+
+
+def create_footer():
+    footer_html = """
+    <style>
+        .sticky-footer {
+            position: fixed;
+            bottom: 0;
+            left: 0;
+            width: 100%;
+            background: white;
+            padding: 10px;
+            box-shadow: 0 -2px 10px rgba(0,0,0,0.1);
+            z-index: 1000;
+        }
+        .content-wrap {
+            padding-bottom: 60px; /* Footer height + extra spacing */
+        }
+    </style>
+    
+    <div class="sticky-footer">
+        <div style="text-align: center; font-size: 14px;">
+            Created by <a href="https://vlai.work" target="_blank" style="color: #007BFF; text-decoration: none;">VLAI</a>
+            • AI VIETNAM
+        </div>
+    </div>
+    """
+    return gr.HTML(footer_html)
+
+
 def get_step_footnote_content(step_log: MemoryStep, step_name: str) -> str:
     """Get a footnote string for a step log with duration and token information"""
     step_footnote = f"**{step_name}**"
@@ -37,111 +117,89 @@ def get_step_footnote_content(step_log: MemoryStep, step_name: str) -> 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_content = f"""<span style="color: #bbbbc2; font-size: 12px;">{step_footnote}</span> """
+    step_footnote_content = f""" < span style = "color: #bbbbc2; font-size: 12px;" > {step_footnote} < /span > """
     return step_footnote_content
 
 
-def pull_messages_from_step(
-    step_log: MemoryStep,
-):
+def pull_messages_from_step(step_log: MemoryStep):
     """Extract ChatMessage objects from agent steps with proper nesting"""
     if not _is_package_available("gradio"):
         raise ModuleNotFoundError(
             "Please install 'gradio' extra to use the GradioUI: `pip install 'smolagents[gradio]'`"
         )
-    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 "Step"
         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
+        if hasattr(step_log, "model_output") and step_log.model_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 = 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, create a parent message
-        if hasattr(step_log, "tool_calls") and step_log.tool_calls is not None:
+        if hasattr(step_log, "tool_calls") and step_log.tool_calls:
             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()
+            content = str(args.get("answer", args)) if isinstance(
+                args, dict) else 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()
+                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(
+            yield gr.ChatMessage(
                 role="assistant",
                 content=content,
                 metadata={
                     "title": f"🛠️ Used tool {first_tool_call.name}",
-                    "id": parent_id,
+                    "id": f"call_{len(step_log.tool_calls)}",
                     "status": "done",
                 },
             )
-            yield parent_message_tool
-
-        # Display execution logs 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"```bash\n{log_content}\n",
-                    metadata={"title": "📝 Execution Logs", "status": "done"},
-                )
 
-        # Display any errors
-        if hasattr(step_log, "error") and step_log.error is not None:
+        if hasattr(step_log, "observations") and step_log.observations and step_log.observations.strip():
+            log_content = re.sub(r"^Execution logs:\s*",
+                                 "", step_log.observations.strip())
+            yield gr.ChatMessage(
+                role="assistant",
+                content=f"```bash\n{log_content}\n```",
+                metadata={"title": "📝 Execution Logs", "status": "done"},
+            )
+
+        if hasattr(step_log, "error") and step_log.error:
             yield gr.ChatMessage(
                 role="assistant",
                 content=str(step_log.error),
                 metadata={"title": "💥 Error", "status": "done"},
             )
 
-        # Update parent message metadata to done status without yielding a new message
         if getattr(step_log, "observations_images", []):
             for image in step_log.observations_images:
                 path_image = AgentImage(image).to_string()
                 yield gr.ChatMessage(
                     role="assistant",
-                    content={"path": path_image, "mime_type": f"image/{path_image.split('.')[-1]}"},
+                    content={"path": path_image,
+                             "mime_type": f"image/{path_image.split('.')[-1]}"},
                     metadata={"title": "🖼️ Output Image", "status": "done"},
                 )
 
-        # Handle standalone errors but not from tool calls
-        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"})
-
         yield gr.ChatMessage(role="assistant", content=get_step_footnote_content(step_log, step_number))
         yield gr.ChatMessage(role="assistant", content="-----", metadata={"status": "done"})
 
     elif isinstance(step_log, PlanningStep):
         yield gr.ChatMessage(role="assistant", content="**Planning step**")
         yield gr.ChatMessage(role="assistant", content=step_log.plan)
-        yield gr.ChatMessage(role="assistant", content=get_step_footnote_content(step_log, "Planning step"))
+        yield gr.ChatMessage(
+            role="assistant",
+            content=get_step_footnote_content(step_log, "Planning step")
+        )
         yield gr.ChatMessage(role="assistant", content="-----", metadata={"status": "done"})
 
     elif isinstance(step_log, FinalAnswerStep):
@@ -154,16 +212,20 @@ def pull_messages_from_step(
         elif isinstance(final_answer, AgentImage):
             yield gr.ChatMessage(
                 role="assistant",
-                content={"path": final_answer.to_string(), "mime_type": "image/png"},
+                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"},
+                content={"path": final_answer.to_string(),
+                         "mime_type": "audio/wav"},
             )
         else:
-            yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
-
+            yield gr.ChatMessage(
+                role="assistant",
+                content=f"**Final answer:** {str(final_answer)}"
+            )
     else:
         raise ValueError(f"Unsupported step type: {type(step_log)}")
 
@@ -179,7 +241,6 @@ def stream_to_gradio(
     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 getattr(agent.model, "last_input_token_count", None) is not None:
             total_input_tokens += agent.model.last_input_token_count
             total_output_tokens += agent.model.last_output_token_count
@@ -187,9 +248,7 @@ def stream_to_gradio(
                 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,
-        ):
+        for message in pull_messages_from_step(step_log):
             yield message
 
 
@@ -205,14 +264,12 @@ class GradioUI:
         self.file_upload_folder = file_upload_folder
         self.name = getattr(agent, "name") or "Agent interface"
         self.description = getattr(agent, "description", None)
-        if self.file_upload_folder is not None:
-            if not os.path.exists(file_upload_folder):
-                os.mkdir(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, session_state):
         import gradio as gr
 
-        # Get the agent type from the template agent
         if "agent" not in session_state:
             session_state["agent"] = self.agent
 
@@ -226,14 +283,11 @@ class GradioUI:
 
             yield messages
         except Exception as e:
-            print(f"Error in interaction: {str(e)}")
-            messages.append(gr.ChatMessage(role="assistant", content=f"Error: {str(e)}"))
+            messages.append(gr.ChatMessage(
+                role="assistant", content=f"Error: {str(e)}"))
             yield messages
 
     def upload_file(self, file, file_uploads_log, allowed_file_types=None):
-        """
-        Handle file uploads, default allowed types are .pdf, .docx, and .txt
-        """
         import gradio as gr
 
         if file is None:
@@ -246,14 +300,9 @@ class GradioUI:
         if file_ext 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
-
-        # Save the uploaded file to the specified folder
-        file_path = os.path.join(self.file_upload_folder, os.path.basename(sanitized_name))
+        sanitized_name = re.sub(r"[^\w\-.]", "_", original_name)
+        file_path = os.path.join(self.file_upload_folder, sanitized_name)
         shutil.copy(file.name, file_path)
 
         return gr.Textbox(f"File uploaded: {file_path}", visible=True), file_uploads_log + [file_path]
@@ -264,9 +313,8 @@ class GradioUI:
         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 ""
+                f"\nYou have been provided with these files: {file_uploads_log}"
+                if file_uploads_log else ""
             ),
             "",
             gr.Button(interactive=False),
@@ -278,92 +326,94 @@ class GradioUI:
     def create_app(self):
         import gradio as gr
 
-        with gr.Blocks(theme="ocean", fill_height=True) as demo:
-            # Add session state to store session-specific data
+        with gr.Blocks(css=CUSTOM_CSS, theme="ocean", fill_height=True) as demo:
+            create_header()
+
             session_state = gr.State({})
             stored_messages = gr.State([])
             file_uploads_log = gr.State([])
 
-            with gr.Sidebar(width=400):
-                gr.Markdown(
-                    f"# {self.name.replace('_', ' ').capitalize()}"
-                    "\n### **AIO2024 Module 10**"
-                )
-                
-                with gr.Row():
-                    logo_base64 = image_to_base64("static/aivn_logo.png")
-                    gr.HTML(f"""
-                    <img src="data:image/png;base64,{logo_base64}" 
-                        alt="Logo" 
-                        style="display: block; margin: auto; height: 120px; width: auto; margin-bottom: 20px;">
-                    """)
-                gr.Markdown(f"\n\n>{self.description}" if self.description else "")
-                    
-                with gr.Group():
-                    gr.Markdown("**Your request**", container=True)
+            # Main content area: Chat + Input
+            with gr.Row(equal_height=True, variant="panel", elem_classes="content-wrap"):
+                # Column for chat and input
+                with gr.Column(scale=3):
+                    # Input area moved here
+                    # gr.Markdown("**Your request**")
                     text_input = gr.Textbox(
-                        lines=3,
-                        label="Chat Message",
-                        container=False,
-                        placeholder="Enter your prompt here and press Shift+Enter or press the button",
+                        lines=2,
+                        label="Your request",
+                        placeholder="Enter your prompt here and press Shift+Enter or the button",
                     )
-                    submit_btn = gr.Button("Submit", variant="primary")
-
-                # 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],
+                    submit_btn = gr.Button(
+                        "Submit", variant="primary", elem_classes="full-width-btn"
                     )
 
+                    # Chatbot
+                    chatbot = gr.Chatbot(
+                        label="Agent",
+                        type="messages",
+                        avatar_images=(
+                            None,
+                            "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
+                        ),
+                        resizeable=True,
+                        scale=1,
+                    )
 
-            # Main chat interface
-            chatbot = gr.Chatbot(
-                label="Agent",
-                type="messages",
-                avatar_images=(
+                # Optional: Column for file uploads
+                if self.file_upload_folder is not None:
+                    with gr.Column(scale=1):
+                        gr.Markdown("**Upload Files**")
+                        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],
+                        )
+
+                # Wiring interactions
+                text_input.submit(
+                    self.log_user_message,
+                    [text_input, file_uploads_log],
+                    [stored_messages, text_input, submit_btn],
+                ).then(
+                    self.interact_with_agent,
+                    [stored_messages, chatbot, session_state],
+                    [chatbot],
+                ).then(
+                    lambda: (
+                        gr.update(value="", interactive=True,
+                                  placeholder="Enter your prompt here and press Shift+Enter or the button"),
+                        gr.update(interactive=True),
+                    ),
                     None,
-                    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
-                ),
-                resizeable=True,
-                scale=1,
-            )
+                    [text_input, submit_btn],
+                )
 
-            # Set up event handlers
-            text_input.submit(
-                self.log_user_message,
-                [text_input, file_uploads_log],
-                [stored_messages, text_input, submit_btn],
-            ).then(self.interact_with_agent, [stored_messages, chatbot, session_state], [chatbot]).then(
-                lambda: (
-                    gr.Textbox(
-                        interactive=True, placeholder="Enter your prompt here and press Shift+Enter or the button"
+                submit_btn.click(
+                    self.log_user_message,
+                    [text_input, file_uploads_log],
+                    [stored_messages, text_input, submit_btn],
+                ).then(
+                    self.interact_with_agent,
+                    [stored_messages, chatbot, session_state],
+                    [chatbot],
+                ).then(
+                    lambda: (
+                        gr.update(value="", interactive=True,
+                                  placeholder="Enter your prompt here and press Shift+Enter or the button"),
+                        gr.update(interactive=True),
                     ),
-                    gr.Button(interactive=True),
-                ),
-                None,
-                [text_input, submit_btn],
-            )
+                    None,
+                    [text_input, submit_btn],
+                )
 
-            submit_btn.click(
-                self.log_user_message,
-                [text_input, file_uploads_log],
-                [stored_messages, text_input, submit_btn],
-            ).then(self.interact_with_agent, [stored_messages, chatbot, session_state], [chatbot]).then(
-                lambda: (
-                    gr.Textbox(
-                        interactive=True, placeholder="Enter your prompt here and press Shift+Enter or the button"
-                    ),
-                    gr.Button(interactive=True),
-                ),
-                None,
-                [text_input, submit_btn],
-            )
+            create_footer()
 
         return demo
 
 
-__all__ = ["stream_to_gradio", "GradioUI"]
+__all__ = ["stream_to_gradio", "GradioUI"]