# 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. # The MIT License # Copyright (c) 2025 Albert Murienne # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN # THE SOFTWARE. import re from typing import Generator from smolagents.agent_types import AgentAudio, AgentImage, AgentText from smolagents.agents import MultiStepAgent, PlanningStep from smolagents.memory import ActionStep, FinalAnswerStep from smolagents.models import ChatMessageStreamDelta, MessageRole, agglomerate_stream_deltas FINAL_ANSWER_TAG = "Final answer:" def get_step_footnote_content(step_log: ActionStep | PlanningStep, step_name: str) -> str: """Get a footnote string for a step log with duration and token information""" step_footnote = f"**{step_name}**" if step_log.token_usage is not None: step_footnote += f" | Input tokens: {step_log.token_usage.input_tokens:,} | Output tokens: {step_log.token_usage.output_tokens:,}" step_footnote += f" | Duration: {round(float(step_log.timing.duration), 2)}s" if step_log.timing.duration else "" step_footnote_content = f"""{step_footnote} """ return step_footnote_content def _clean_model_output(model_output: str) -> str: """ Clean up model output by removing trailing tags and extra backticks. Args: model_output (`str`): Raw model output. Returns: `str`: Cleaned model output. """ if not model_output: return "" model_output = model_output.strip() # Remove any trailing and extra backticks, handling multiple possible formats model_output = re.sub(r"```\s*", "```", model_output) # handles ``` model_output = re.sub(r"\s*```", "```", model_output) # handles ``` model_output = re.sub(r"```\s*\n\s*", "```", model_output) # handles ```\n return model_output.strip() def _format_code_content(content: str) -> str: """ Format code content as Python code block if it's not already formatted. Args: content (`str`): Code content to format. Returns: `str`: Code content formatted as a Python code block. """ content = content.strip() # Remove existing code blocks and end_code tags content = re.sub(r"```.*?\n", "", content) content = re.sub(r"\s*\s*", "", content) content = content.strip() # Add Python code block formatting if not already present if not content.startswith("```python"): content = f"```python\n{content}\n```" return content def _process_action_step(step_log: ActionStep, skip_model_outputs: bool = False) -> Generator: """ Process an [`ActionStep`] and yield appropriate Gradio ChatMessage objects. Args: step_log ([`ActionStep`]): ActionStep to process. skip_model_outputs (`bool`): Whether to skip model outputs. Yields: `gradio.ChatMessage`: Gradio ChatMessages representing the action step. """ import gradio as gr # Output the step number step_number = f"Step {step_log.step_number}" if not skip_model_outputs: yield gr.ChatMessage(role=MessageRole.ASSISTANT, content=f"**{step_number}**", metadata={"status": "done"}) # First yield the thought/reasoning from the LLM if not skip_model_outputs and getattr(step_log, "model_output", ""): model_output = _clean_model_output(step_log.model_output) yield gr.ChatMessage(role=MessageRole.ASSISTANT, content=model_output, metadata={"status": "done"}) # For tool calls, create a parent message if getattr(step_log, "tool_calls", []): first_tool_call = step_log.tool_calls[0] used_code = first_tool_call.name == "python_interpreter" # Process 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() # Format code content if needed if used_code: content = _format_code_content(content) # Create the tool call message parent_message_tool = gr.ChatMessage( role=MessageRole.ASSISTANT, content=content, metadata={ "title": f"🛠️ Used tool {first_tool_call.name}", "status": "done", }, ) yield parent_message_tool # Display execution logs if they exist if getattr(step_log, "observations", "") 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=MessageRole.ASSISTANT, content=f"```bash\n{log_content}\n", metadata={"title": "📝 Execution Logs", "status": "done"}, ) # Display any images in observations if getattr(step_log, "observations_images", []): for image in step_log.observations_images: path_image = AgentImage(image).to_string() yield gr.ChatMessage( role=MessageRole.ASSISTANT, content={"path": path_image, "mime_type": f"image/{path_image.split('.')[-1]}"}, metadata={"title": "🖼️ Output Image", "status": "done"}, ) # Handle errors if getattr(step_log, "error", None): yield gr.ChatMessage( role=MessageRole.ASSISTANT, content=str(step_log.error), metadata={"title": "💥 Error", "status": "done"} ) # Add step footnote and separator yield gr.ChatMessage( role=MessageRole.ASSISTANT, content=get_step_footnote_content(step_log, step_number), metadata={"status": "done"}, ) yield gr.ChatMessage(role=MessageRole.ASSISTANT, content="-----", metadata={"status": "done"}) def _process_planning_step(step_log: PlanningStep, skip_model_outputs: bool = False) -> Generator: """ Process a [`PlanningStep`] and yield appropriate gradio.ChatMessage objects. Args: step_log ([`PlanningStep`]): PlanningStep to process. Yields: `gradio.ChatMessage`: Gradio ChatMessages representing the planning step. """ import gradio as gr if not skip_model_outputs: yield gr.ChatMessage(role=MessageRole.ASSISTANT, content="**Planning step**", metadata={"status": "done"}) yield gr.ChatMessage(role=MessageRole.ASSISTANT, content=step_log.plan, metadata={"status": "done"}) yield gr.ChatMessage( role=MessageRole.ASSISTANT, content=get_step_footnote_content(step_log, "Planning step"), metadata={"status": "done"}, ) yield gr.ChatMessage(role=MessageRole.ASSISTANT, content="-----", metadata={"status": "done"}) def _process_final_answer_step(step_log: FinalAnswerStep) -> Generator: """ Process a [`FinalAnswerStep`] and yield appropriate gradio.ChatMessage objects. Args: step_log ([`FinalAnswerStep`]): FinalAnswerStep to process. Yields: `gradio.ChatMessage`: Gradio ChatMessages representing the final answer. """ import gradio as gr final_answer = step_log.output if isinstance(final_answer, AgentText): yield gr.ChatMessage( role=MessageRole.ASSISTANT, content=f"**{FINAL_ANSWER_TAG}**\n{final_answer.to_string()}\n", metadata={"status": "done"}, ) elif isinstance(final_answer, AgentImage): yield gr.ChatMessage( role=MessageRole.ASSISTANT, content={"path": final_answer.to_string(), "mime_type": "image/png"}, metadata={"status": "done"}, ) elif isinstance(final_answer, AgentAudio): yield gr.ChatMessage( role=MessageRole.ASSISTANT, content={"path": final_answer.to_string(), "mime_type": "audio/wav"}, metadata={"status": "done"}, ) else: yield gr.ChatMessage( role=MessageRole.ASSISTANT, content=f"**{FINAL_ANSWER_TAG}** {str(final_answer)}", metadata={"status": "done"} ) def pull_messages_from_step(step_log: ActionStep | PlanningStep | FinalAnswerStep, skip_model_outputs: bool = False): """Extract Gradio ChatMessage objects from agent steps with proper nesting. Args: step_log: The step log to display as gr.ChatMessage objects. skip_model_outputs: If True, skip the model outputs when creating the gr.ChatMessage objects: This is used for instance when streaming model outputs have already been displayed. """ if isinstance(step_log, ActionStep): yield from _process_action_step(step_log, skip_model_outputs) elif isinstance(step_log, PlanningStep): yield from _process_planning_step(step_log, skip_model_outputs) elif isinstance(step_log, FinalAnswerStep): yield from _process_final_answer_step(step_log) else: raise ValueError(f"Unsupported step type: {type(step_log)}") def stream_to_gradio( agent, task: str, additional_args: dict | None = None, ) -> Generator: """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages.""" accumulated_events: list[ChatMessageStreamDelta] = [] for event in agent.run(task, additional_args=additional_args): if isinstance(event, ActionStep | PlanningStep | FinalAnswerStep): for message in pull_messages_from_step( event, # If we're streaming model outputs, no need to display them twice skip_model_outputs=getattr(agent, "stream_outputs", False), ): yield message accumulated_events = [] elif isinstance(event, ChatMessageStreamDelta): accumulated_events.append(event) text = agglomerate_stream_deltas(accumulated_events).render_as_markdown() yield text class AgentUI: """ Gradio interface for interacting with a [`MultiStepAgent`]. This class provides a web interface to interact with the agent in real-time, allowing users to submit prompts, and receive responses in a chat-like format. It can reset the agent's memory at the start of each interaction if desired. It uses the [`gradio.Chatbot`] component to display the conversation history. This class requires the `gradio` extra to be installed: `pip install 'smolagents[gradio]'`. Args: agent ([`MultiStepAgent`]): The agent to interact with. """ def __init__(self, agent: MultiStepAgent): self.agent = agent self.description = getattr(agent, "description", None) def set_advanced_mode(self, enabled: bool): """ Configure the agent to enable/disable advanced mode. """ self.agent.enable_advanced_mode(enabled) def interact_with_agent(self, prompt: str, verbose_messages: list, quiet_messages: list): """ Interacts with the agent and streams results into two separate histories: - verbose_messages: full reasoning stream (Chatterbox) - quiet_messages: only user prompt + final answer (Quiet) Quiet is enhanced with pending "Step N..." indicators only (no generic thinking text). """ import gradio as gr try: # Append the user message to both histories (quiet keeps the user query) user_msg = gr.ChatMessage(role="user", content=prompt, metadata={"status": "done"}) verbose_messages.append(user_msg) quiet_messages.append(user_msg) # yield initial state to update UI immediately yield verbose_messages, quiet_messages quiet_pending_idx = None for msg in stream_to_gradio(self.agent, task=prompt): # Full gr.ChatMessage object (from steps) — append to verbose always if isinstance(msg, gr.ChatMessage): # Mark last verbose pending -> done if needed and append if verbose_messages and verbose_messages[-1].metadata.get("status") == "pending": verbose_messages[-1].metadata["status"] = "done" verbose_messages[-1].content = msg.content else: verbose_messages.append(msg) content_text = msg.content if isinstance(msg.content, str) else "" # Detect final answer messages and append to quiet # HACK : FinalAnswerStep messages are produced by _process_final_answer_step and use FINAL_ANSWER_TAG if FINAL_ANSWER_TAG in content_text: # Remove everything before and including the FINAL_ANSWER_TAG label (and any leading/trailing whitespace/newlines) answer_only = re.sub( rf"(?s)^.*?\*\*{FINAL_ANSWER_TAG}\*\*\s*[\n]*", # (?s) allows . to match newlines "", content_text, flags=re.IGNORECASE, ) final_msg = gr.ChatMessage(role=MessageRole.ASSISTANT, content=answer_only, metadata={"status": "done"}) if quiet_pending_idx is not None: quiet_messages[quiet_pending_idx] = final_msg quiet_pending_idx = None else: quiet_messages.append(final_msg) else: # Look for "Step " pattern match = re.search(r"\bStep\s*(\d+)\b", content_text, re.IGNORECASE) if match: step_num = match.group(1) pending_text = f"⏳ Step {step_num}..." if quiet_pending_idx is None: quiet_messages.append( gr.ChatMessage( role=MessageRole.ASSISTANT, content=pending_text, metadata={"status": "pending"}, ) ) quiet_pending_idx = len(quiet_messages) - 1 else: quiet_messages[quiet_pending_idx].content = pending_text elif isinstance(msg, str): text = msg.replace("<", r"\<").replace(">", r"\>") if verbose_messages and verbose_messages[-1].metadata.get("status") == "pending": verbose_messages[-1].content = text else: verbose_messages.append( gr.ChatMessage(role=MessageRole.ASSISTANT, content=text, metadata={"status": "pending"}) ) yield verbose_messages, quiet_messages # final yield to ensure both UIs are up-to-date yield verbose_messages, quiet_messages except Exception as e: # ensure UIs don't hang if something failed yield verbose_messages, quiet_messages raise gr.Error(f"Error in interaction: {str(e)}") def clear_history(self): """ Clear the chat history and reset the agent's memory. """ self.agent.reset() return [], [] def disable_query(self, text_input): """ Disable the text input and submit button while the agent is processing. """ import gradio as gr return ( text_input, gr.Textbox( value="", placeholder="Wait for answer completion before submitting a new prompt...", interactive=False ), gr.Button(interactive=False), ) def enable_query(self): """ Enable the text input and submit button after the agent has finished processing. """ import gradio as gr return ( gr.Textbox( interactive=True, placeholder="Enter your prompt here and press Shift+Enter or the button" ), gr.Button(interactive=True), ) def launch(self, share: bool = True, **kwargs): """ Launch the Gradio app with the agent interface. Args: share (`bool`, defaults to `True`): Whether to share the app publicly. **kwargs: Additional keyword arguments to pass to the Gradio launch method. """ self.create_app().launch(debug=True, share=share, **kwargs) def get_tavily_credits(self): """ Fetch the Tavily credits. """ return self.agent.get_search_credits() def get_advanced_mode(self) -> bool: """ Return the agent's current advanced_mode flag for initializing the checkbox on page load. """ return getattr(self.agent, "advanced_mode", False) def create_app(self): import gradio as gr # some nice thmes available here: https://huggingface.co/spaces/gradio/theme-gallery with gr.Blocks(theme="JohnSmith9982/small_and_pretty", fill_height=True) as agent: # Set up states to hold the session information stored_query = gr.State("") # current user query stored_messages_verbose = gr.State([]) # full reasoning history stored_messages_quiet = gr.State([]) # only user + final answer with gr.Sidebar(): gr.Markdown( "# SmolAlbert 🤖" ) with gr.Group(): gr.Markdown("**Your request**", container=True) text_input = gr.Textbox( lines=3, label="Chat Message", container=False, placeholder="Enter your prompt here and press Shift+Enter or press the button", ) submit_btn = gr.Button("Submit", variant="primary") # Advanced search mode checkbox advanced_checkbox = gr.Checkbox( label="Advanced search mode", value=getattr(self.agent, "advanced_mode", False), info="Toggle advanced search behavior for the agent (x2 search credits).", container=True, ) # call agent configuration when checkbox changes advanced_checkbox.change(self.set_advanced_mode, advanced_checkbox, None) # ensure the checkbox reflects the current agent state each time a page/session loads agent.load(self.get_advanced_mode, None, advanced_checkbox) tavily_credits = gr.Textbox( label="Tavily Credits", value=self.get_tavily_credits(), interactive=False, container=True, ) gr.HTML( "

Powered by smolagents

" ) with gr.Tab("Quiet", scale=1): quiet_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, latex_delimiters=[ {"left": r"$$", "right": r"$$", "display": True}, {"left": r"$", "right": r"$", "display": False}, {"left": r"\[", "right": r"\]", "display": True}, {"left": r"\(", "right": r"\)", "display": False}, ], ) with gr.Tab("Chatterbox", scale=1): verbose_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, latex_delimiters=[ {"left": r"$$", "right": r"$$", "display": True}, {"left": r"$", "right": r"$", "display": False}, {"left": r"\[", "right": r"\]", "display": True}, {"left": r"\(", "right": r"\)", "display": False}, ], ) # Main input handlers: call interact_with_agent(prompt, verbose_state, quiet_state) text_input.submit( self.disable_query, text_input, [stored_query, text_input, submit_btn] ).then( self.interact_with_agent, [stored_query, stored_messages_verbose, stored_messages_quiet], [verbose_chatbot, quiet_chatbot], ).then( self.get_tavily_credits, None, tavily_credits, ).then( self.enable_query, None, [text_input, submit_btn], ) submit_btn.click( self.disable_query, text_input, [stored_query, text_input, submit_btn] ).then( self.interact_with_agent, [stored_query, stored_messages_verbose, stored_messages_quiet], [verbose_chatbot, quiet_chatbot], ).then( self.get_tavily_credits, None, tavily_credits, ).then( self.enable_query, None, [text_input, submit_btn], ) # bind clears to both chat components so agent memory is reset quiet_chatbot.clear(self.clear_history, inputs=None, outputs=[stored_messages_verbose, stored_messages_quiet]) verbose_chatbot.clear(self.clear_history, inputs=None, outputs=[stored_messages_verbose, stored_messages_quiet]) return agent