prompt_builder.py

import json


class PromptBuilder:
    def __init__(self, prompt=None, prompt_type="chat_completion"):
        if prompt is None:
            prompt = []

        self.prompt = prompt
        self.prompt_type = prompt_type

    @staticmethod
    def load_from_file(file_path):
        with open(file_path, "r") as fp:
            content = fp.read()

            try:
                # Attempt to parse as JSON list (chat completion format)
                prompts = json.loads(content)
                if isinstance(prompts, list):
                    return PromptBuilder(prompts, "chat_completion")
                
                raise ValueError("Unsupported prompt type!")
            
            except json.JSONDecodeError:
                return PromptBuilder(content, "completion")


    def save_to_file(self, file_path):
        with open(file_path, "w") as fp:
            if self.prompt_type == "chat_completion":
                json.dump(self.prompt, fp, ensure_ascii=False, indent=4)
            elif prompt_type == "completion":
                fp.write(self.prompt)


    def append(self, prompt):
        assert isinstance(prompt, list), "Prompt must be a list!"
        self.prompt.append(prompt)

    def format(self, **kwargs):
        result = None

        if self.prompt_type == "chat_completion":
            result = []
            for prompt in self.prompt:
                result.append({
                    "role": prompt["role"],
                    "content": prompt["content"].format(**kwargs)
                })
        elif self.prompt_type == "completion":
            result = self.prompt.format(**kwargs)

        return result

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        if self.prompt_type == "chat_completion":
            return json.dumps(self.prompt, indent=4)
        elif self.prompt_type == "completion":
            return self.prompt


def system(content):
    return {"role": "system", "content": content}


def user(content):
    return {"role": "user", "content": content}


def assistant(content):
    return {"role": "assistant", "content": content}