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

class TritonPythonModel:
    def initialize(self, args):
        self.logger = pb_utils.Logger

        self.model_name = args["model_name"]
        self.model_config = json.loads(args["model_config"])

        self.logger.log_info(f"'{self.model_name}' 모델 초기화 완료")


    def execute(self, requests):
        responses = []

        for request in requests:
            # 랜덤 문자열 생성
            random_string = ''.join(
                random.choice(string.ascii_letters + string.digits) for _ in range(16)
            )
            self.logger.log_info(f"OUTPUT 출력:\n{random_string}")

            # Triton이 요구하는 shape [1] STRING 출력
            output_tensor = pb_utils.Tensor(
                "OUTPUT",
                np.array([random_string.encode("utf-8")], dtype=np.object_)
            )

            responses.append(pb_utils.InferenceResponse(
                output_tensors=[output_tensor]
            ))

        return responses   

    def finalize(self):
        pass