# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/01_base.ipynb (unless otherwise specified).

__all__ = ['SEPARATOR_TOKEN', 'Message', 'Conversation', 'Config', 'Prompt']

# Cell
from dataclasses import dataclass
from typing import Optional, List

# Cell
SEPARATOR_TOKEN = '<|endoftext|>'

# Cell
@dataclass(frozen=True)
class Message:
    user: str
    text: Optional[str] = None

    def render(self):
        result = self.user + ':'
        if self.text is not None:
            result += ' ' + self.text
        return result

# Cell
@dataclass
class Conversation:
    messages: List[Message]

    def prepend(self, message: Message):
        self.messages.insert(0, message)
        return self

    def render(self):
        return f'\n{SEPARATOR_TOKEN}'.join(
            [message.render() for message in self.messages]
        )

# Cell
@dataclass(frozen=True)
class Config:
    name: str
    instructions: str
    example_conversations: List[Conversation]

# Cell
@dataclass(frozen=True)
class Prompt:
    header: Message
    examples: List[Conversation]
    convo: Conversation

    def render(self):
        return f'\n{SEPARATOR_TOKEN}'.join(
            [self.header.render()]
            + [Message('System', 'Example conversations:').render()]
            + [conversation.render() for conversation in self.examples]
            + [Message('System', 'Current conversation:').render()]
            + [self.convo.render()],
        )