35 lines
599 B
Python
35 lines
599 B
Python
|
from enum import StrEnum, auto
|
||
|
|
|
||
|
|
from pydantic import BaseModel, ConfigDict
|
||
|
|
|
||
|
|
|
||
|
|
class Role(StrEnum):
|
||
|
|
"""Role of a user in the chat."""
|
||
|
|
|
||
|
|
ASSISTANT = auto()
|
||
|
|
SYSTEM = auto()
|
||
|
|
USER = auto()
|
||
|
|
|
||
|
|
|
||
|
|
class Message(BaseModel):
|
||
|
|
"""Message in the conversation."""
|
||
|
|
|
||
|
|
role: Role
|
||
|
|
content: str
|
||
|
|
|
||
|
|
model_config: ConfigDict = ConfigDict(
|
||
|
|
frozen=True,
|
||
|
|
use_enum_values=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class Conversation(BaseModel):
|
||
|
|
"""Conversation in the chat."""
|
||
|
|
|
||
|
|
messages: list[Message]
|
||
|
|
|
||
|
|
model_config: ConfigDict = ConfigDict(
|
||
|
|
frozen=False,
|
||
|
|
use_enum_values=True,
|
||
|
|
)
|