KaniumCogs/reginaldCog/llm_clients/openai_data_models.py

from dataclasses import dataclass, field
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 | None = field(default=None)
    file_id: str | None = 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(frozen=True)
class Prompt:
    model: str
    input: list[Message]


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