KaniumCogs/reginaldCog/openai_client/models.py

from dataclasses import dataclass, field, asdict, fields
from abc import ABC


# region Content classes
@dataclass(frozen=True)
class Content(ABC):
    type: str = field(init=False, default='')


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


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


@dataclass(frozen=True)
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(frozen=True)
class UrlCitation(Content):
    type: str = field(init=False, default='url_citation')
    # To be done


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


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


@dataclass(frozen=True)
class Message:
    role: str
    content: list[Content]


@dataclass
class Prompt:
    model: str
    input: list[Message]

    def add_message(self, message: Message):
        self.input.append(message)


@dataclass(frozen=True)
class Response:
    output: list[Message]


def filter_fields(data: dict, filter_by_class: dataclass) -> dict:
    return {
        i_key: i_value
        for i_key, i_value in data.items()
        if i_key in tuple(f.name for f in fields(filter_by_class))
    }


if __name__ == '__main__':
    from openai import OpenAI
    client = OpenAI()

    test_system_message = Message(role='system', content=[InputText('Talk like an Italian mafia boss.')])
    test_user_message = Message(role='user', content=[InputText('Hi! How are you?')])
    test_prompt = Prompt(model='gpt-4.1-mini', input=[test_system_message, test_user_message])
    print(asdict(test_prompt))
    response_raw = client.responses.create(**asdict(test_prompt))
    response = Response(**filter_fields(response_raw.to_dict(), Response))
    print(asdict(response))
    test_prompt.add_message(response.output[0])

    test_user_message_2 = Message(role='user', content=[
        InputText('Can you tell me what is on this picture?'),
        InputImage(image_url='https://upload.wikimedia.org/wikipedia/commons/f/f4/Piet_Mondriaan_-_Sinaasappelen_%28authentiek%29_-_A97_-_Piet_Mondrian%2C_catalogue_raisonn%C3%A9.jpg')
    ])
    test_prompt.add_message(test_user_message_2)
    response_raw = client.responses.create(**asdict(test_prompt))
    response = Response(**filter_fields(response_raw.to_dict(), Response))
    test_prompt.add_message(response.output[0])
    print(asdict(test_prompt))