KaniumCogs/reginaldCog/openai/content_builder.py

from dataclasses import dataclass, field, fields


@dataclass
class Content:
    type: str = field(init=False, default='')

    def to_dict(self) -> dict:
        result = {'type': self.type}

        for i_field in fields(self):
            if i_field.name == 'type':
                continue

            value = getattr(self, i_field.name)

            if value is None:
                continue

            if isinstance(value, Content):
                result[i_field.name] = value.to_dict()
            elif isinstance(value, list):
                result[i_field.name] = [
                    i_content.to_dict() if isinstance(i_content, Content) else i_content
                    for i_content in value if i_content is not None
                ]
            else:
                result[i_field.name] = value

        return result


@dataclass
class InputText(Content):
    type: str = field(init=False, default='input_text')
    text: str


@dataclass
class OutputText(Content):
    type: str = field(init=False, default='output_text')
    text: str


@dataclass
class InputImage(Content):
    type: str = field(init=False, default='input_image')
    image_url: str = field(default=None)
    file_id: str = field(default=None)

    def __post_init__(self):
        if self.image_url is None and self.file_id is None:
            raise ValueError("Either 'image_url' or 'file_id' must be provided.")


@dataclass
class UrlCitation(Content):
    pass


@dataclass
class FunctionCall(Content):
    type: str = field(init=False, default='function_call')
    id: str
    call_id: str
    name: str
    arguments: dict


@dataclass
class FunctionCallOutput(Content):
    type: str = field(init=False, default='function_call_output')
    call_id: str
    output: str


if __name__ == '__main__':
    test_input_text = InputText(text='What\'s on this image?')
    print(test_input_text.to_dict())
    test_input_image = InputImage(image_url='https://upload.wikimedia.org/wikipedia/commons/7/72/Skansen_Bobrka_2.jpg')
    print(test_input_image.to_dict())