DeepSpaceSearch / my_ui.py
2stacks's picture
Integrate MCP servers and refactor app architecture
ab0e2df verified
raw
history blame
5.03 kB
from functools import partial
import gradio as gr
from smolagents import GradioUI
class CustomGradioUI(GradioUI):
"""Custom GradioUI that allows customization of the smolagents default interface."""
def __init__(
self, agent, file_upload_folder=None, reset_agent_memory=False, allowed_file_types=None
):
super().__init__(agent, file_upload_folder, reset_agent_memory)
self.allowed_file_types = allowed_file_types or [".pdf", ".docx", ".txt"]
def create_app(self):
"""Override create_app to use custom allowed_file_types."""
# Call parent's create_app but we need to rebuild it with our custom upload handler
with gr.Blocks(theme="ocean", fill_height=True) as demo:
session_state = gr.State({})
stored_messages = gr.State([])
file_uploads_log = gr.State([])
with gr.Sidebar():
gr.Markdown(
f"# {self.name.replace('_', ' ').capitalize()}"
"\n> This web ui allows you to interact with a `smolagents` agent "
"that can use tools and execute steps to complete tasks."
+ (
f"\n\n**Agent description:**\n{self.description}"
if self.description
else ""
)
)
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")
# 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
)
# Use partial to bind allowed_file_types
upload_handler = partial(
self.upload_file, allowed_file_types=self.allowed_file_types
)
upload_file.change(
upload_handler,
[upload_file, file_uploads_log],
[upload_status, file_uploads_log],
)
gr.HTML(
"<br><br><h4><center>Powered by <a target='_blank' href='https://github.com/huggingface/smolagents'><b>smolagents</b></a></center></h4>"
)
# Main chat interface
chatbot = gr.Chatbot(
label="Agent",
type="messages",
avatar_images=(
None,
"https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/smolagents/mascot_smol.png",
),
resizable=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},
],
)
# 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",
),
gr.Button(interactive=True),
),
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],
)
return demo