AgentCoreService/pqai_agent/toolkit/message_toolkit.py

from typing import List, Dict

from pqai_agent.logging import logger
from pqai_agent.toolkit.base import BaseToolkit
from pqai_agent.toolkit.function_tool import FunctionTool
from pqai_agent.toolkit.tool_registry import register_toolkit


@register_toolkit
class MessageToolkit(BaseToolkit):
    def __init__(self):
        super().__init__()

    def message_notify_user(self, message: str) -> str:
        """Sends a message to the user.

        Args:
            message (str): The message to send.
        Returns:
            str: A confirmation message.
        """

        logger.info(f"Message to user: {message}")
        return 'success'

    def output_multimodal_message(self, message: Dict[str, str]) -> str:
        """Outputs a multimodal message to the user.
        Message schema:
        {
            "type": "text|image|gif|video|mini_program|link",
            "content": "text message content or url of the media",
            "title": "only needed if type in: video, link, mini_program",
            "cover_url": "cover image url, only needed if type in: mini_program",
            "desc": "description, optional if type in: link"
        }
        if message type is image, gif, video, link or mini_program, the content should be a URL.

        Args:
            message (Dict[str, str]): The message to output.
        Returns:
            str: A confirmation message.
        """
        msg_type = message.get("type", "")
        if msg_type not in ["text", "image", "gif", "video", "mini_program", "link"]:
            return f"Invalid message type: {msg_type}"
        if msg_type in ("video", "mini_program", "link") and "title" not in message:
            return f"Title is required for [{msg_type}] messages."
        if msg_type in ("mini_program", ) and "cover_url" not in message:
            return f"Cover image URL is required for [{msg_type}] messages."
        # if msg_type in ("link", ) and "desc" not in message:
        #     return f"Description is required for [link] messages."
        logger.info(f"Multimodal message to user: {message}")
        return 'success'


    def get_tools(self):
        return [FunctionTool(self.message_notify_user),
                FunctionTool(self.output_multimodal_message)]