web_ui.py

import os
import pprint
import re
from typing import List, Optional, Union

from qwen_agent import Agent, MultiAgentHub
from qwen_agent.agents.user_agent import PENDING_USER_INPUT
from qwen_agent.gui.gradio_utils import format_cover_html
from qwen_agent.gui.utils import convert_fncall_to_text, convert_history_to_chatbot, get_avatar_image
from qwen_agent.llm.schema import CONTENT, FILE, IMAGE, NAME, ROLE, USER, Message
from qwen_agent.log import logger
from qwen_agent.utils.utils import print_traceback
import spaces


class WebUI:
    """A chatbot application for EV advantages."""

    def __init__(self, agent: Union[Agent, MultiAgentHub, List[Agent]], chatbot_config: Optional[dict] = None):
        """
        Initialize the chatbot.
        Args:
            agent: The agent or a list of agents (Assistant or MultiAgentHub).
            chatbot_config: Configuration for the chatbot, including placeholder text and prompts.
        """
        chatbot_config = chatbot_config or {}

        if isinstance(agent, MultiAgentHub):
            self.agent_list = [agent for agent in agent.nonuser_agents]
            self.agent_hub = agent
        elif isinstance(agent, list):
            self.agent_list = agent
            self.agent_hub = None
        else:
            self.agent_list = [agent]
            self.agent_hub = None

        user_name = chatbot_config.get('user.name', 'User')
        self.user_config = {
            'name': user_name,
            'avatar': chatbot_config.get(
                'user.avatar',
                get_avatar_image(user_name),
            ),
        }

        self.agent_config_list = [{
            'name': agent.name,
            'avatar': chatbot_config.get(
                'agent.avatar',
                self._resolve_logo_path('logo.jpeg')
            ),
            'description': agent.description or "I'm here to help with EV advantages.",
        } for agent in self.agent_list]

        self.input_placeholder = chatbot_config.get('input.placeholder', "Start your EV-related conversation here...")
        self.prompt_suggestions = chatbot_config.get('prompt.suggestions', [])
        self.verbose = chatbot_config.get('verbose', False)

    def _resolve_logo_path(self, filename: str) -> str:
        """
        Resolve the avatar/logo path, falling back to a default if not found.
        """
        path = os.path.join('/home/user/app', filename)
        if not os.path.exists(path):
            logger.warning(f"Avatar file '{filename}' not found. Using placeholder.")
            return os.path.join('/home/user/app', 'default_logo.jpeg')
        return path

    @spaces.GPU
    def run(self,
            messages: List[Message] = None,
            share: bool = False,
            server_name: str = None,
            server_port: int = None,
            concurrency_limit: int = 80,
            enable_mention: bool = False,
            **kwargs):
        """
        Run the chatbot interface.
        Args:
            messages: The initial chat history (if any).
        """
        self.run_kwargs = kwargs

        from qwen_agent.gui.gradio import gr, mgr

        customTheme = gr.themes.Default(
            primary_hue=gr.themes.utils.colors.blue,
            radius_size=gr.themes.utils.sizes.radius_none,
        )

        with gr.Blocks(css=os.path.join(os.path.dirname(__file__), 'assets/appBot.css'), theme=customTheme) as demo:
            history = gr.State([])

            with gr.Row(elem_classes='container'):
                with gr.Column(scale=4):
                    chatbot = mgr.Chatbot(
                        value=convert_history_to_chatbot(messages=messages),
                        avatar_images=[self.user_config, self.agent_config_list],
                        height=850,
                        avatar_image_width=80,
                        flushing=False,
                        show_copy_button=True,
                        latex_delimiters=[{'left': '\\(', 'right': '\\)', 'display': True}]
                    )

                    input_box = mgr.MultimodalInput(placeholder=self.input_placeholder, upload_button_props=dict(visible=False))

                with gr.Column(scale=1):
                    if len(self.agent_list) > 1:
                        agent_selector = gr.Dropdown(
                            [(agent.name, i) for i, agent in enumerate(self.agent_list)],
                            label='Select Agent',
                            info='Choose an Agent',
                            value=0,
                            interactive=True,
                        )

                    agent_info_block = self._create_agent_info_block()

                    if self.prompt_suggestions:
                        gr.Examples(
                            label='Suggested Conversations',
                            examples=self.prompt_suggestions,
                            inputs=[input_box],
                        )

                input_promise = input_box.submit(
                    fn=self.add_text,
                    inputs=[input_box, chatbot, history],
                    outputs=[input_box, chatbot, history],
                    queue=False,
                )

                input_promise.then(
                    self.agent_run,
                    [chatbot, history],
                    [chatbot, history],
                ).then(self.flushed, None, [input_box])

            demo.queue(default_concurrency_limit=concurrency_limit).launch(
                share=share, server_name=server_name, server_port=server_port
            )

    @spaces.GPU
    def add_text(self, _input, _chatbot, _history):
        """
        Process user input and update the chat history.
        """
        from qwen_agent.gui.gradio import gr
        if _input.text == "/clear":
            _chatbot = []
            _history.clear()
            yield gr.update(interactive=False, value=""), _chatbot, _history
            return

        _history.append({
            ROLE: USER,
            CONTENT: [{'text': _input.text}],
        })

        if self.user_config[NAME]:
            _history[-1][NAME] = self.user_config[NAME]

        _chatbot.append([_input, None])

        yield gr.update(interactive=False, value=None), _chatbot, _history

    @spaces.GPU
    def agent_run(self, _chatbot, _history, _agent_selector=None):
        """
        Run the agent and generate responses.
        """
        if not _history:
            yield _chatbot, _history
            return

        if self.verbose:
            logger.info('Agent run input:\n' + pprint.pformat(_history, indent=2))

        agent_runner = self.agent_list[_agent_selector or 0]
        responses = []
        for responses in agent_runner.run(_history, **self.run_kwargs):
            if not responses:
                continue

            display_responses = convert_fncall_to_text(responses)
            if not display_responses:
                continue

            _chatbot[-1][1] = display_responses[-1][CONTENT]

            yield _chatbot, _history

        if responses:
            _history.extend([res for res in responses])

    def flushed(self):
        from qwen_agent.gui.gradio import gr
        return gr.update(interactive=True)

    def _create_agent_info_block(self, agent_index=0):
        from qwen_agent.gui.gradio import gr

        agent_config = self.agent_config_list[agent_index]
        return gr.HTML(
            format_cover_html(
                bot_name=agent_config['name'],
                bot_description=agent_config['description'],
                bot_avatar=agent_config['avatar'],
            )
        )