import logging
from towhee.operator import PyOperator
from towhee import triton_client


log = logging.getLogger()


class Client(PyOperator):
    """
    Triton client class.
    
    Args:
        url: str
            IP address and https port for the triton server, such as '<host>:<port>', '127.0.0.1:8001'.
        model_name: str
            The name of the model to run inference, defaults to 'pipline'.
    """
    def __init__(self, url, model_name='pipeline'):
        super().__init__()
        self._client = triton_client(url, model_name)

    def __call__(self, *data):
        """
        A one line summary of this function.

        Args:
            data:
                This data to your triton server.
        """
        return self._client(*data)

    def batch(self, data_list):
        return self._client.batch(data_list)