AgentCoreService/pqai_agent/toolkit/message_notifier.py

from typing import List, Dict

from pqai_agent.logging_service 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 MessageNotifier(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",
            "content": "message content",
            "title": "only needed if type is video or mini_program",
            "cover_image": "only needed if type is mini_program",
        }
        if message type is image, gif, video or mini_program, the content should be a URL.

        Args:
            message (Dict[str, str]): The message to output.
        Returns:
            str: A confirmation message.
        """
        if message["type"] not in ["text", "image", "gif", "video", "mini_program"]:
            return f"Invalid message type: {message['type']}"
        if message["type"] in ("video", "mini_program") and "title" not in message:
            return "Title is required for video or mini_program messages."
        if message["type"] == "mini_program" and "cover_image" not in message:
            return "Cover image is required for mini_program 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)]