import triton_python_backend_utils as pb_utils
import numpy as np
import random
import string
import json

class TritonPythonModel:
    """
    Triton Python Model 클래스.
    """

def initialize(self, args):
    print("TritonPythonModel: initialize() called.")
    self.model_config = json.loads(args['model_config'])
    
    # 출력 설정에서 데이터 타입 정보를 가져옴
    output_config = pb_utils.get_output_config_by_name(
        self.model_config, "OUTPUT")
    
    # Triton 데이터 타입 문자열을 NumPy 데이터 타입으로 직접 변환
    # 'BYTES'는 np.object_ 타입에 해당함
    if output_config['data_type'] == 'TYPE_STRING':
        self.output_dtype = np.object_
    else:
        # 다른 데이터 타입에 대한 처리 로직 추가 가능
        self.output_dtype = pb_utils.triton_string_to_np_dtype(output_config['data_type'])

    def execute(self, requests):
        """
        클라이언트의 요청이 들어올 때마다 호출됩니다.
        """
        responses = []

        for request in requests:
            # 입력 데이터 가져오기 (이 예제에서는 입력 데이터를 사용하지 않음)
            in_0 = pb_utils.get_input_tensor_by_name(request, "INPUT")
            
            # 무작위 문자열 생성
            random_string = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(16))

            output = random_string

            # 생성된 텍스트를 Triton 출력 텐서로 변환합니다.
            output_tensor = pb_utils.Tensor("OUTPUT", np.array(output.encode('utf-8'), dtype=np.bytes_))

            # 응답 객체를 생성하고 출력 텐서를 추가합니다.
            responses.append(pb_utils.InferenceResponse(output_tensors=[output_tensor]))

        return responses

    def finalize(self):
        """
        모델이 언로드될 때 호출됩니다.
        """
        print("TritonPythonModel: finalize() called.")