src/tools/user_profile_toolkit.py

import asyncio
import chainlit as cl
from phi.tools import Toolkit
from phi.utils.log import logger

from src.libs.rpc_client import rpc_call


class UserProfileToolkit(Toolkit):
    def __init__(self):
        super().__init__(name="user_profile_toolkit")

        # Registering methods to make them accessible via the toolkit
        self.register(self.get_user_name)
        self.register(self.get_user_email)
        self.register(self.get_user_picture)
        self.register(self.get_all_user_info)
        self.register(self.update_user_info)
        self.register(self.update_user_name)
        self.register(self.update_user_email)
        self.register(self.update_user_picture)
        self.register(self.get_user_id)

        # fetching user email
        self.email = self.get_user_email()

    @cl.on_chat_start
    def get_user_info(self, info_type: str) -> str:
        """
        Fetches user information from the Chainlit user session.

        Args:
        info_type (str): The type of information to fetch ('name', 'email', 'picture').

        Returns:
        str: The requested user information.

        Example:
        >>> get_user_info('name')
        """
        logger.info(f"Fetching user info: {info_type}")
        try:
            # Fetch user session data
            user_session = cl.user_session.get("user").metadata
            if not user_session:
                raise ValueError("User session not found")

            response = user_session[info_type]
            return f"{response}"
        except Exception as e:
            logger.warning(f"Failed to get user info: {e}")
            return f"Error: {e}"

    def get_all_user_info(self) -> str:
        """
        Fetches all user information from the Chainlit user session.

        Returns:
        dict: A dictionary containing all user information.

        Example:
        >>> get_all_user_info()
        """
        logger.info("Fetching all user info")
        try:
            # Fetch user session data
            user_session = cl.user_session.get("user").metadata
            if not user_session:
                raise ValueError("User session not found")

            response = user_session
            return f"{response}"
        except Exception as e:
            logger.warning(f"Failed to get all user info: {e}")
            # return {"error": str(e)}
            return f"Error: {e}"

    def update_user_info(self, info_type: str, value: str) -> str:
        """
        Updates user information in the Chainlit user session.

        Args:
        info_type (str): The type of information to update ('name', 'email', 'picture').
        value (str): The new value to set.

        Returns:
        str: Confirmation message.

        Example:
        >>> update_user_info('name', 'Jane Doe')
        """
        logger.info(f"Updating user info: {info_type} to {value}")
        try:
            # Fetch user session data
            user_session = cl.user_session.get("user").metadata
            if not user_session:
                raise ValueError("User session not found")

            # Update the user session data
            user_session[info_type] = value
            cl.user_session.set("user", user_session)
            return f"{info_type} updated to {value}"
        except Exception as e:
            logger.warning(f"Failed to update user info: {e}")
            return f"Error: {e}"

    def get_user_name(self) -> str:
        """
        Fetches the user's name from the Chainlit user session.

        Returns:
        str: The user's name.

        Example:
        >>> get_user_name()
        """
        response = self.get_user_info('name')
        return f"{response}"

    def get_user_email(self) -> str:
        """
        Fetches the user's email from the Chainlit user session.

        Returns:
        str: The user's email.

        Example:
        >>> get_user_email()
        """
        response = self.get_user_info('email')
        return f"{response}"

    def get_user_picture(self) -> str:
        """
        Fetches the user's picture URL from the Chainlit user session.

        Returns:
        str: The user's picture URL.

        Example:
        >>> get_user_picture()
        """
        response = self.get_user_info('picture')
        return f"{response}"

    def update_user_name(self, name: str) -> str:
        """
        Updates the user's name in the Chainlit user session.

        Args:
        name (str): The new name to set.

        Returns:
        str: Confirmation message.

        Example:
        >>> update_user_name('Jane Doe')
        """
        response = self.update_user_info('name', name)
        return f"{response}"

    def update_user_email(self, email: str) -> str:
        """
        Updates the user's email in the Chainlit user session.

        Args:
        email (str): The new email to set.

        Returns:
        str: Confirmation message.

        Example:
        >>> update_user_email('jane.doe@chatxbt.com')
        """
        response = self.update_user_info('email', email)
        return f"{response}"

    def update_user_picture(self, picture: str) -> str:
        """
        Updates the user's picture URL in the Chainlit user session.

        Args:
        picture (str): The new picture URL to set.

        Returns:
        str: Confirmation message.

        Example:
        >>> update_user_picture('https://example.com/new_picture.jpg')
        """
        response = self.update_user_info('picture', picture)
        return f"{response}"
    
    def get_user_id(self, user_email: str) -> str:
        """
        Fetches the ChatXBT ID of the user associated with the given email.

        Args:
        user_email (str): The email address of the user whose ChatXBT ID needs to be fetched.

        Returns:
        str]: A string containing the ChatXBT ID of the user. If the user is not found, an empty string is returned.

        Example:
        >>> get_user_id('jane.doe@chatxbt.com')
        """
        logger.info("Fetching user ChatXBT ID")

        params = {
            "userEmail": user_email.lower()
        }
        response = asyncio.run(rpc_call(method_name="getUserId", params=params))
        return f"{response}"